Frequency Table of Tree Distance
(tree/frequency_table_of_tree_distance.cpp)

• View this file on GitHub
• Last update: 2026-03-23 22:54:37+09:00

説明

木上の距離ごとの頂点対数を求める。 res[d] に距離 d の unordered pair 数を入れて返す。 計算量は $O(N \log^2 N)$。

できること

• vector<long long> frequency_table_of_tree_distance(const vector<vector<int>>& g) 距離頻度表を返す。res[0] = 0、res[d] は距離 d の頂点対数

使い方

隣接リストを渡して呼ぶ。

実装上の補足

• 重心分解で各重心を通る頂点対だけを数える
• 深さ頻度の畳み込みを使い、同一子部分木内の対を引いて重複を除く

Depends on

• 任意MOD畳み込み(3 NTT + Garner) (math/fft.cpp)

Verified with

• test/yosupo_frequency_table_of_tree_distance.test.cpp

Code

#define FIRIEXP_LIBRARY_FFT_EXACT_ONLY
#include "../math/fft.cpp"
#undef FIRIEXP_LIBRARY_FFT_EXACT_ONLY

vector<long long> frequency_table_of_tree_distance(const vector<vector<int>> &g) {
    using ArbitraryConvolution::square_u64;

    int n = (int)g.size();
    if (n == 0) return {};

    vector<int> sz(n);
    vector<char> used(n, 0);
    vector<long long> ordered(n);

    auto dfs_size = [&](auto &&self, int v, int p) -> int {
        sz[v] = 1;
        for (auto &&to : g[v]) {
            if (to == p || used[to]) continue;
            sz[v] += self(self, to, v);
        }
        return sz[v];
    };

    auto find_centroid = [&](auto &&self, int v, int p, int half) -> int {
        for (auto &&to : g[v]) {
            if (to == p || used[to]) continue;
            if (sz[to] > half) return self(self, to, v, half);
        }
        return v;
    };

    auto collect = [&](auto &&self, int v, int p, int d, vector<unsigned int> &cnt) -> void {
        if ((int)cnt.size() <= d) cnt.resize(d + 1);
        ++cnt[d];
        for (auto &&to : g[v]) {
            if (to == p || used[to]) continue;
            self(self, to, v, d + 1, cnt);
        }
    };

    auto add_convolution = [&](const vector<unsigned long long> &conv, int sign) {
        int lim = min(n, (int)conv.size());
        for (int d = 0; d < lim; ++d) {
            ordered[d] += sign * (long long)conv[d];
        }
    };

    auto solve = [&](auto &&self, int start) -> void {
        int centroid = find_centroid(find_centroid, start, -1, dfs_size(dfs_size, start, -1) / 2);
        used[centroid] = 1;

        vector<unsigned int> total(1, 1);
        for (auto &&to : g[centroid]) {
            if (used[to]) continue;
            vector<unsigned int> cnt(1, 0);
            collect(collect, to, centroid, 1, cnt);
            add_convolution(square_u64(cnt), -1);
            if (total.size() < cnt.size()) total.resize(cnt.size());
            for (int d = 1; d < (int)cnt.size(); ++d) total[d] += cnt[d];
        }
        add_convolution(square_u64(total), 1);

        for (auto &&to : g[centroid]) {
            if (!used[to]) self(self, to);
        }
    };

    solve(solve, 0);

    vector<long long> ans(n);
    for (int d = 1; d < n; ++d) ans[d] = ordered[d] / 2;
    return ans;
}

/**
 * @brief Frequency Table of Tree Distance
 */

Traceback (most recent call last):
  File "/opt/hostedtoolcache/Python/3.10.20/x64/lib/python3.10/site-packages/onlinejudge_verify/documentation/build.py", line 71, in _render_source_code_stat
    bundled_code = language.bundle(stat.path, basedir=basedir, options={'include_paths': [basedir]}).decode()
  File "/opt/hostedtoolcache/Python/3.10.20/x64/lib/python3.10/site-packages/onlinejudge_verify/languages/cplusplus.py", line 187, in bundle
    bundler.update(path)
  File "/opt/hostedtoolcache/Python/3.10.20/x64/lib/python3.10/site-packages/onlinejudge_verify/languages/cplusplus_bundle.py", line 401, in update
    self.update(self._resolve(pathlib.Path(included), included_from=path))
  File "/opt/hostedtoolcache/Python/3.10.20/x64/lib/python3.10/site-packages/onlinejudge_verify/languages/cplusplus_bundle.py", line 355, in update
    raise BundleErrorAt(path, i + 1, "found codes out of include guard")
onlinejudge_verify.languages.cplusplus_bundle.BundleErrorAt: math/fft.cpp: line 346: found codes out of include guard

Back to top page