test/yosupo/frequency_table_of_tree_distance.test.cpp

• View this file on GitHub
• Last update: 2024-01-08 00:27:17+09:00
• Problem: https://judge.yosupo.jp/problem/frequency_table_of_tree_distance

Depends on

• Fast Fourier Transform (convolution/fft.hpp)
• Centroid Decomposition (tree/centroid_decomposition.hpp)

Code

#define PROBLEM "https://judge.yosupo.jp/problem/frequency_table_of_tree_distance"

#include "../../tree/centroid_decomposition.hpp"
#include "../../convolution/fft.hpp"

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define rep(i, s, t) for (int i = (int)(s); i < (int)(t); ++i)
#define revrep(i, t, s) for (int i = (int)(t)-1; i >= (int)(s); --i)
#define all(x) begin(x), end(x)
template <typename T> bool chmax(T& a, const T& b) { return a < b ? (a = b, 1) : 0; }
template <typename T> bool chmin(T& a, const T& b) { return a > b ? (a = b, 1) : 0; }

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout << fixed << setprecision(15);

    int N;
    cin >> N;
    vector<vector<int>> G(N);
    rep(i,0,N-1) {
        int a, b;
        cin >> a >> b;
        G[a].push_back(b);
        G[b].push_back(a);
    }

    vector<int> level, sz, par;
    tie(level, sz, par) = centroid_decomposition(G);

    vector<ll> ans(N);
    vector<double> cnt;

    auto dfs = [&](auto& dfs, int v, int p, int d, int l) -> void {
        if (cnt.size() <= d) cnt.push_back(0);
        ++cnt[d];
        for (int c : G[v]) {
            if (c != p && level[c] > l) {
                dfs(dfs, c, v, d+1, l);
            }
        }
    };

    auto calc = [&](int s) {
        vector<double> res(sz[s]);
        cnt.assign(1, 0);
        for (int c : G[s]) {
            if (level[c] > level[s]) {
                dfs(dfs, c, s, 1, level[s]);
            }
        }
        auto conv = convolution(cnt, cnt);
        rep(d,0,cnt.size()) {
            res[d] += cnt[d];
        }
        rep(d,0,min(conv.size(), res.size())) {
            res[d] += conv[d]/2;
        }
        for (int c : G[s]) {
            if (level[c] > level[s]) {
                cnt.assign(1, 0);
                dfs(dfs, c, s, 1, level[s]);
                auto conv = convolution(cnt, cnt);
                rep(d,0,min(conv.size(), res.size())) {
                    res[d] -= conv[d]/2;
                }
            }
        }
        rep(d,0,sz[s]) {
            ans[d] += (ll) round(res[d]);
        }
    };

    rep(i,0,N) calc(i);
    rep(d,1,N) cout << ans[d] << (d < N-1 ? " " : "\n");
}

#line 1 "test/yosupo/frequency_table_of_tree_distance.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/frequency_table_of_tree_distance"

#line 2 "tree/centroid_decomposition.hpp"
#include <tuple>
#include <vector>

std::tuple<std::vector<int>, std::vector<int>, std::vector<int>>
centroid_decomposition(const std::vector<std::vector<int>>& G) {
    const int N = G.size();
    std::vector<int> sz(N), level(N, -1), sz_comp(N), par(N);

    auto dfs_size = [&](auto& dfs_size, int v, int p) -> int {
        sz[v] = 1;
        for (int c : G[v]) {
            if (c != p && level[c] == -1) sz[v] += dfs_size(dfs_size, c, v);
        }
        return sz[v];
    };

    auto dfs_centroid = [&](auto& dfs_centroid, int v, int p, int n) -> int {
        for (int c : G[v]) {
            if (c != p && level[c] == -1 && sz[c] > n / 2)
                return dfs_centroid(dfs_centroid, c, v, n);
        }
        return v;
    };

    auto decompose = [&](auto& decompose, int v, int k, int p) -> void {
        int n = dfs_size(dfs_size, v, -1);
        int s = dfs_centroid(dfs_centroid, v, -1, n);
        level[s] = k;
        sz_comp[s] = n;
        par[s] = p;
        for (int c : G[s]) {
            if (level[c] == -1) decompose(decompose, c, k + 1, s);
        }
    };

    decompose(decompose, 0, 0, -1);
    return {level, sz_comp, par};
}
#line 2 "convolution/fft.hpp"
#include <bit>
#include <complex>
#include <numbers>
#line 6 "convolution/fft.hpp"

void fft(std::vector<std::complex<double>>& a) {
    const int n = a.size();
    for (int m = n; m > 1; m >>= 1) {
        double ang = 2.0 * std::numbers::pi / m;
        std::complex<double> omega(cos(ang), sin(ang));
        for (int s = 0; s < n / m; ++s) {
            std::complex<double> w(1, 0);
            for (int i = 0; i < m / 2; ++i) {
                auto l = a[s * m + i];
                auto r = a[s * m + i + m / 2];
                a[s * m + i] = l + r;
                a[s * m + i + m / 2] = (l - r) * w;
                w *= omega;
            }
        }
    }
}

void ifft(std::vector<std::complex<double>>& a) {
    const int n = a.size();
    for (int m = 2; m <= n; m <<= 1) {
        double ang = -2.0 * std::numbers::pi / m;
        std::complex<double> omega(cos(ang), sin(ang));
        for (int s = 0; s < n / m; ++s) {
            std::complex<double> w(1, 0);
            for (int i = 0; i < m / 2; ++i) {
                auto l = a[s * m + i];
                auto r = a[s * m + i + m / 2] * w;
                a[s * m + i] = l + r;
                a[s * m + i + m / 2] = l - r;
                w *= omega;
            }
        }
    }
}

template <typename T>
std::vector<double> convolution(const std::vector<T>& a,
                                const std::vector<T>& b) {
    const int size = a.size() + b.size() - 1;
    const int n = std::bit_ceil((unsigned int)size);
    std::vector<std::complex<double>> na(a.begin(), a.end()),
        nb(b.begin(), b.end());
    na.resize(n);
    nb.resize(n);
    fft(na);
    fft(nb);
    for (int i = 0; i < n; ++i) na[i] *= nb[i];
    ifft(na);
    std::vector<double> ret(size);
    for (int i = 0; i < size; ++i) ret[i] = na[i].real() / n;
    return ret;
}
#line 5 "test/yosupo/frequency_table_of_tree_distance.test.cpp"

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define rep(i, s, t) for (int i = (int)(s); i < (int)(t); ++i)
#define revrep(i, t, s) for (int i = (int)(t)-1; i >= (int)(s); --i)
#define all(x) begin(x), end(x)
template <typename T> bool chmax(T& a, const T& b) { return a < b ? (a = b, 1) : 0; }
template <typename T> bool chmin(T& a, const T& b) { return a > b ? (a = b, 1) : 0; }

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout << fixed << setprecision(15);

    int N;
    cin >> N;
    vector<vector<int>> G(N);
    rep(i,0,N-1) {
        int a, b;
        cin >> a >> b;
        G[a].push_back(b);
        G[b].push_back(a);
    }

    vector<int> level, sz, par;
    tie(level, sz, par) = centroid_decomposition(G);

    vector<ll> ans(N);
    vector<double> cnt;

    auto dfs = [&](auto& dfs, int v, int p, int d, int l) -> void {
        if (cnt.size() <= d) cnt.push_back(0);
        ++cnt[d];
        for (int c : G[v]) {
            if (c != p && level[c] > l) {
                dfs(dfs, c, v, d+1, l);
            }
        }
    };

    auto calc = [&](int s) {
        vector<double> res(sz[s]);
        cnt.assign(1, 0);
        for (int c : G[s]) {
            if (level[c] > level[s]) {
                dfs(dfs, c, s, 1, level[s]);
            }
        }
        auto conv = convolution(cnt, cnt);
        rep(d,0,cnt.size()) {
            res[d] += cnt[d];
        }
        rep(d,0,min(conv.size(), res.size())) {
            res[d] += conv[d]/2;
        }
        for (int c : G[s]) {
            if (level[c] > level[s]) {
                cnt.assign(1, 0);
                dfs(dfs, c, s, 1, level[s]);
                auto conv = convolution(cnt, cnt);
                rep(d,0,min(conv.size(), res.size())) {
                    res[d] -= conv[d]/2;
                }
            }
        }
        rep(d,0,sz[s]) {
            ans[d] += (ll) round(res[d]);
        }
    };

    rep(i,0,N) calc(i);
    rep(d,1,N) cout << ans[d] << (d < N-1 ? " " : "\n");
}

Back to top page