Test/Manual/yahoo_procon2018_final_c.test.cpp

• View this file on GitHub
• Last update: 2024-12-06 16:40:42+09:00
• Problem: https://onlinejudge.u-aizu.ac.jp/courses/lesson/2/ITP1/1/ITP1_1_A
• Link: https://atcoder.jp/contests/yahoo-procon2018-final-open/submissions/60475048

Depends on

• 木の重心 (Src/Graph/Tree/Centroid.hpp)
• ioまわりの設定 (Src/Template/IOSetting.hpp)
• 標準データ型のエイリアス (Src/Template/TypeAlias.hpp)

Code

#define PROBLEM "https://onlinejudge.u-aizu.ac.jp/courses/lesson/2/ITP1/1/ITP1_1_A"

/*
 * 「みんなのプロコン 2018」決勝 オープンコンテスト - C 木の問題
 * https://atcoder.jp/contests/yahoo-procon2018-final-open/submissions/60475048
 */

#include "../../Src/Template/IOSetting.hpp"
#include "../../Src/Graph/Tree/Centroid.hpp"

#include <iostream>
#include <utility>
#include <vector>

using namespace zawa;

void solve() {
    int N, Q;
    std::cin >> N >> Q;
    std::vector<std::vector<int>> g(N);
    for (int _{} ; _ < N - 1 ; _++) {
        int A, B;
        std::cin >> A >> B;
        A--; B--;
        g[A].push_back(B);
        g[B].push_back(A);
    }
    std::vector<std::vector<std::pair<int, int>>> query(N);
    for (int i{} ; i < Q ; i++) {
        int v, k;
        std::cin >> v >> k;
        v--;
        query[v].push_back(std::make_pair(k, i));
    }
    Centroid C(std::move(g));
    std::vector<int> dep(N);
    auto calc_dep{[&](auto dfs, int v, int p, int d) -> void {
        dep[v] = d;
        for (auto x : C[v]) if (!C.isRemoved(x) and (int)x != p) {
            dfs(dfs, x, v, d + 1);
        }
    }};
    std::vector<int> ans(Q);
    auto dfs{[&](auto dfs, int v) -> void {
        v = C.rooting(v);
        calc_dep(calc_dep, v, -1, 0);
        C.remove(v);
        std::vector<std::vector<int>> comp;
        for (auto x : C.adjlist(v)) {
            comp.push_back(C.component(x));
        }
        int max_d{};
        for (const auto& com : comp) {
            for (auto x : com) max_d = std::max(max_d, dep[x]);
        }
        std::vector<int> cnt(max_d + 1);
        cnt[0]++;
        for (const auto& com : comp) {
            for (auto x : com) cnt[dep[x]]++;
        }
        for (auto [k, i] : query[v]) {
            if (k <= max_d) {
                ans[i] += cnt[k];
            }
        }
        for (const auto& com : comp) {
            for (auto x : com) cnt[dep[x]]--;
            for (auto x : com) {
                for (auto [k, i] : query[x]) {
                    int d{k - dep[x]};
                    if (0 <= d and d <= max_d) ans[i] += cnt[d];
                }
            }
            for (auto x : com) cnt[dep[x]]++;
        }
        for (auto x : C.adjlist(v)) {
            dfs(dfs, x);
        }
    }};
    dfs(dfs, 0);
    for (int i{} ; i < Q ; i++) {
        std::cout << ans[i] << '\n';
    }
}

int main() {
#ifdef ATCODER
    SetFastIO();
    solve();
#else
    std::cout << "Hello World" << '\n';
#endif
}

#line 1 "Test/Manual/yahoo_procon2018_final_c.test.cpp"
#define PROBLEM "https://onlinejudge.u-aizu.ac.jp/courses/lesson/2/ITP1/1/ITP1_1_A"

/*
 * 「みんなのプロコン 2018」決勝 オープンコンテスト - C 木の問題
 * https://atcoder.jp/contests/yahoo-procon2018-final-open/submissions/60475048
 */

#line 2 "Src/Template/IOSetting.hpp"

#line 2 "Src/Template/TypeAlias.hpp"

#include <cstdint>
#include <cstddef>

namespace zawa {

using i16 = std::int16_t;
using i32 = std::int32_t;
using i64 = std::int64_t;
using i128 = __int128_t;

using u8 = std::uint8_t;
using u16 = std::uint16_t;
using u32 = std::uint32_t;
using u64 = std::uint64_t;

using usize = std::size_t;

} // namespace zawa
#line 4 "Src/Template/IOSetting.hpp"

#include <iostream>
#include <iomanip>

namespace zawa {

void SetFastIO() {
    std::cin.tie(nullptr)->sync_with_stdio(false);
}

void SetPrecision(u32 dig) {
    std::cout << std::fixed << std::setprecision(dig);
}

} // namespace zawa
#line 2 "Src/Graph/Tree/Centroid.hpp"

#line 4 "Src/Graph/Tree/Centroid.hpp"

#include <cassert>
#include <utility>
#include <vector>

namespace zawa {

template <class V>
class Centroid {
public:

    Centroid() = default;

    Centroid(const std::vector<std::vector<V>>& T) : T_{T}, size_(T_.size(), usize{1}) {}
    Centroid(std::vector<std::vector<V>>&& T) : T_{std::move(T)}, size_(T_.size(), usize{1}) {}

    inline usize size() const noexcept {
        return T_.size();
    }

    inline usize size(V v) const noexcept {
        assert(v < (V)size());
        return size_[v];
    }

    bool isRemoved(V v) const noexcept {
        assert(v < (V)size());
        return size_[v] == 0u;
    }

    void remove(V v) noexcept {
        assert(v < (V)size());
        size_[v] = 0u;
    }

    const std::vector<V>& operator[](V v) const noexcept {
        assert(v < (V)size());
        return T_[v];
    }

    // @response: centroid of component which v belongs
    V rooting(V v) {
        assert(v < (V)size());
        assert(!isRemoved(v));
        usize all{dfsSize(v, INVALID)};
        V par{INVALID};
        bool fn{false};
        while (!fn) {
            fn = true;
            for (V x : T_[v]) {
                if (x == par or isRemoved(x) or usize{2} * size_[x] <= all) {
                    continue;
                }
                fn = false;
                par = v;
                v = x;
                break;
            }
        }
        return v;
    }

    std::vector<V> component(V v) const {
        assert(v < (V)size());
        assert(!isRemoved(v));
        std::vector<V> res;
        dfsComponent(v, INVALID, res);
        return res;
    }

    std::vector<V> adjlist(V v) const {
        assert(v < (V)size());
        std::vector<V> res;
        res.reserve(T_[v].size());
        for (auto x : T_[v]) if (!isRemoved(x)) {
            res.emplace_back(x);
        }
        return res;
    }

private:
    static constexpr V INVALID{static_cast<V>(-1)};
    std::vector<std::vector<V>> T_{};
    std::vector<usize> size_{};

    usize dfsSize(V v, V par) {
        size_[v] = 1u;
        for (V x : T_[v]) if (x != par and !isRemoved(x)) {
            size_[v] += dfsSize(x, v);
        }
        return size_[v];
    }

    void dfsComponent(V v, V par, std::vector<V>& comp) const {
        comp.emplace_back(v);
        for (V x : T_[v]) if (x != par and !isRemoved(x)) {
            dfsComponent(x, v, comp);
        }
    }
};

} // namespace zawa
#line 10 "Test/Manual/yahoo_procon2018_final_c.test.cpp"

#line 14 "Test/Manual/yahoo_procon2018_final_c.test.cpp"

using namespace zawa;

void solve() {
    int N, Q;
    std::cin >> N >> Q;
    std::vector<std::vector<int>> g(N);
    for (int _{} ; _ < N - 1 ; _++) {
        int A, B;
        std::cin >> A >> B;
        A--; B--;
        g[A].push_back(B);
        g[B].push_back(A);
    }
    std::vector<std::vector<std::pair<int, int>>> query(N);
    for (int i{} ; i < Q ; i++) {
        int v, k;
        std::cin >> v >> k;
        v--;
        query[v].push_back(std::make_pair(k, i));
    }
    Centroid C(std::move(g));
    std::vector<int> dep(N);
    auto calc_dep{[&](auto dfs, int v, int p, int d) -> void {
        dep[v] = d;
        for (auto x : C[v]) if (!C.isRemoved(x) and (int)x != p) {
            dfs(dfs, x, v, d + 1);
        }
    }};
    std::vector<int> ans(Q);
    auto dfs{[&](auto dfs, int v) -> void {
        v = C.rooting(v);
        calc_dep(calc_dep, v, -1, 0);
        C.remove(v);
        std::vector<std::vector<int>> comp;
        for (auto x : C.adjlist(v)) {
            comp.push_back(C.component(x));
        }
        int max_d{};
        for (const auto& com : comp) {
            for (auto x : com) max_d = std::max(max_d, dep[x]);
        }
        std::vector<int> cnt(max_d + 1);
        cnt[0]++;
        for (const auto& com : comp) {
            for (auto x : com) cnt[dep[x]]++;
        }
        for (auto [k, i] : query[v]) {
            if (k <= max_d) {
                ans[i] += cnt[k];
            }
        }
        for (const auto& com : comp) {
            for (auto x : com) cnt[dep[x]]--;
            for (auto x : com) {
                for (auto [k, i] : query[x]) {
                    int d{k - dep[x]};
                    if (0 <= d and d <= max_d) ans[i] += cnt[d];
                }
            }
            for (auto x : com) cnt[dep[x]]++;
        }
        for (auto x : C.adjlist(v)) {
            dfs(dfs, x);
        }
    }};
    dfs(dfs, 0);
    for (int i{} ; i < Q ; i++) {
        std::cout << ans[i] << '\n';
    }
}

int main() {
#ifdef ATCODER
    SetFastIO();
    solve();
#else
    std::cout << "Hello World" << '\n';
#endif
}

Back to top page