test/verify/yosupo-frequency-table-of-tree-distance.test.cpp

• View this file on GitHub
• Last update: 2026-08-10 15:36:17+00:00
• Problem: https://judge.yosupo.jp/problem/frequency_table_of_tree_distance

Depends on

• Graph Template (グラフテンプレート) (graph/graph-template.hpp)
• Centroid-Decomposition(重心分解) (graph/tree/centroid-decomposition.hpp)
• math/fft/fast-fourier-transform.hpp

Code

// clang-format off
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/frequency_table_of_tree_distance
// clang-format on

#include <cstdint>
#include <functional>
#include <iostream>
#include <queue>
#include <tuple>
#include <vector>

#include "../../graph/tree/centroid-decomposition.hpp"
#include "../../math/fft/fast-fourier-transform.hpp"

using namespace std;

int main() {
  int N;
  cin >> N;
  CentroidDecomposition<int> g(N);
  g.read(N - 1, 0);
  int root = g.build();
  vector<int> used(N);
  vector<int64_t> dist(N);
  function<void(int)> rec = [&](int centroid) {
    used[centroid] = true;
    vector<int> cnt{1};
    for (auto& ch : g.g[centroid]) {
      if (used[ch]) continue;
      vector<int> num;
      queue<tuple<int, int, int> > que;
      que.emplace(ch, centroid, 1);
      while (!que.empty()) {
        int idx, par, dep;
        tie(idx, par, dep) = que.front();
        que.pop();
        if (cnt.size() <= dep) cnt.resize(dep + 1, 0);
        if (num.size() <= dep) num.resize(dep + 1, 0);
        cnt[dep]++;
        num[dep]++;
        for (auto& to : g.g[idx]) {
          if (to == par || used[to]) continue;
          que.emplace(to.to, idx, dep + 1);
        }
      }
      auto ret = FastFourierTransform::multiply(num, num);
      for (int i = 0; i < ret.size(); i++) dist[i] -= ret[i];
    }
    auto ret = FastFourierTransform::multiply(cnt, cnt);
    for (int i = 0; i < ret.size(); i++) dist[i] += ret[i];
    for (auto& to : g.tree.g[centroid]) rec(to);
  };
  rec(root);
  dist.erase(begin(dist));
  for (auto& p : dist) p /= 2;
  for (int i = 0; i < N - 1; ++i) {
    if (i) cout << ' ';
    cout << dist[i];
  }
  cout << '\n';
}

#line 1 "test/verify/yosupo-frequency-table-of-tree-distance.test.cpp"
// clang-format off
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/frequency_table_of_tree_distance
// clang-format on

#include <cstdint>
#include <functional>
#include <iostream>
#include <queue>
#include <tuple>
#include <vector>

#line 2 "graph/tree/centroid-decomposition.hpp"

#line 4 "graph/tree/centroid-decomposition.hpp"

#line 2 "graph/graph-template.hpp"

#include <cstddef>
#line 6 "graph/graph-template.hpp"

template <typename T = int>
struct Edge {
  int from, to;
  T cost;
  int idx;

  Edge() = default;

  Edge(int from, int to, T cost = 1, int idx = -1)
      : from(from), to(to), cost(cost), idx(idx) {}

  operator int() const { return to; }
};

template <typename T = int>
struct Graph {
  std::vector<std::vector<Edge<T> > > g;
  int es;

  Graph() = default;

  explicit Graph(int n) : g(n), es(0) {}

  std::size_t size() const { return g.size(); }

  void add_directed_edge(int from, int to, T cost = 1) {
    g[from].emplace_back(from, to, cost, es++);
  }

  void add_edge(int from, int to, T cost = 1) {
    g[from].emplace_back(from, to, cost, es);
    g[to].emplace_back(to, from, cost, es++);
  }

  void read(int M, int padding = -1, bool weighted = false,
            bool directed = false) {
    for (int i = 0; i < M; i++) {
      int a, b;
      std::cin >> a >> b;
      a += padding;
      b += padding;
      T c = T(1);
      if (weighted) std::cin >> c;
      if (directed)
        add_directed_edge(a, b, c);
      else
        add_edge(a, b, c);
    }
  }

  inline std::vector<Edge<T> >& operator[](const int& k) { return g[k]; }

  inline const std::vector<Edge<T> >& operator[](const int& k) const {
    return g[k];
  }
};

template <typename T = int>
using Edges = std::vector<Edge<T> >;
#line 6 "graph/tree/centroid-decomposition.hpp"

/**
 * @brief Centroid-Decomposition(重心分解)
 */
template <typename T>
struct CentroidDecomposition : Graph<T> {
 public:
  using Graph<T>::Graph;
  using Graph<T>::g;
  Graph<int> tree;

  int build(int t = 0) {
    sub.assign(g.size(), 0);
    v.assign(g.size(), 0);
    tree = Graph<int>(g.size());
    return build_dfs(0);
  }

  explicit CentroidDecomposition(const Graph<T>& g) : Graph<T>(g) {}

 private:
  std::vector<int> sub;
  std::vector<int> v;

  inline int build_dfs(int idx, int par) {
    sub[idx] = 1;
    for (auto& to : g[idx]) {
      if (to == par || v[to]) continue;
      sub[idx] += build_dfs(to, idx);
    }
    return sub[idx];
  }

  inline int search_centroid(int idx, int par, const int mid) {
    for (auto& to : g[idx]) {
      if (to == par || v[to]) continue;
      if (sub[to] > mid) return search_centroid(to, idx, mid);
    }
    return idx;
  }

  inline int build_dfs(int idx) {
    int centroid = search_centroid(idx, -1, build_dfs(idx, -1) / 2);
    v[centroid] = true;
    for (auto& to : g[centroid]) {
      if (!v[to]) tree.add_directed_edge(centroid, build_dfs(to));
    }
    v[centroid] = false;
    return centroid;
  }
};
#line 2 "math/fft/fast-fourier-transform.hpp"

#include <algorithm>
#include <cassert>
#include <cmath>
#line 8 "math/fft/fast-fourier-transform.hpp"

namespace FastFourierTransform {
using real = double;

struct C {
  real x, y;

  C() : x(0), y(0) {}

  C(real x, real y) : x(x), y(y) {}

  inline C operator+(const C& c) const { return C(x + c.x, y + c.y); }

  inline C operator-(const C& c) const { return C(x - c.x, y - c.y); }

  inline C operator*(const C& c) const {
    return C(x * c.x - y * c.y, x * c.y + y * c.x);
  }

  inline C conj() const { return C(x, -y); }
};

const real PI = std::acos(-1);
int base = 1;
std::vector<C> rts = {{0, 0}, {1, 0}};
std::vector<int> rev = {0, 1};

void ensure_base(int nbase) {
  if (nbase <= base) return;
  rev.resize(1 << nbase);
  rts.resize(1 << nbase);
  for (int i = 0; i < (1 << nbase); i++) {
    rev[i] = (rev[i >> 1] >> 1) + ((i & 1) << (nbase - 1));
  }
  while (base < nbase) {
    real angle = PI * 2.0 / (1 << (base + 1));
    for (int i = 1 << (base - 1); i < (1 << base); i++) {
      rts[i << 1] = rts[i];
      real angle_i = angle * (2 * i + 1 - (1 << base));
      rts[(i << 1) + 1] = C(std::cos(angle_i), std::sin(angle_i));
    }
    ++base;
  }
}

void fft(std::vector<C>& a, int n) {
  assert((n & (n - 1)) == 0);
  int zeros = __builtin_ctz(n);
  ensure_base(zeros);
  int shift = base - zeros;
  for (int i = 0; i < n; i++) {
    if (i < (rev[i] >> shift)) {
      std::swap(a[i], a[rev[i] >> shift]);
    }
  }
  for (int k = 1; k < n; k <<= 1) {
    for (int i = 0; i < n; i += 2 * k) {
      for (int j = 0; j < k; j++) {
        C z = a[i + j + k] * rts[j + k];
        a[i + j + k] = a[i + j] - z;
        a[i + j] = a[i + j] + z;
      }
    }
  }
}

std::vector<std::int64_t> multiply(const std::vector<int>& a,
                                   const std::vector<int>& b) {
  int need = (int)a.size() + (int)b.size() - 1;
  int nbase = 1;
  while ((1 << nbase) < need) nbase++;
  ensure_base(nbase);
  int sz = 1 << nbase;
  std::vector<C> fa(sz);
  for (int i = 0; i < sz; i++) {
    int x = (i < (int)a.size() ? a[i] : 0);
    int y = (i < (int)b.size() ? b[i] : 0);
    fa[i] = C(x, y);
  }
  fft(fa, sz);
  C r(0, -0.25 / (sz >> 1)), s(0, 1), t(0.5, 0);
  for (int i = 0; i <= (sz >> 1); i++) {
    int j = (sz - i) & (sz - 1);
    C z = (fa[j] * fa[j] - (fa[i] * fa[i]).conj()) * r;
    fa[j] = (fa[i] * fa[i] - (fa[j] * fa[j]).conj()) * r;
    fa[i] = z;
  }
  for (int i = 0; i < (sz >> 1); i++) {
    C A0 = (fa[i] + fa[i + (sz >> 1)]) * t;
    C A1 = (fa[i] - fa[i + (sz >> 1)]) * t * rts[(sz >> 1) + i];
    fa[i] = A0 + A1 * s;
  }
  fft(fa, sz >> 1);
  std::vector<std::int64_t> ret(need);
  for (int i = 0; i < need; i++) {
    ret[i] = std::llround(i & 1 ? fa[i >> 1].y : fa[i >> 1].x);
  }
  return ret;
}
};  // namespace FastFourierTransform
#line 14 "test/verify/yosupo-frequency-table-of-tree-distance.test.cpp"

using namespace std;

int main() {
  int N;
  cin >> N;
  CentroidDecomposition<int> g(N);
  g.read(N - 1, 0);
  int root = g.build();
  vector<int> used(N);
  vector<int64_t> dist(N);
  function<void(int)> rec = [&](int centroid) {
    used[centroid] = true;
    vector<int> cnt{1};
    for (auto& ch : g.g[centroid]) {
      if (used[ch]) continue;
      vector<int> num;
      queue<tuple<int, int, int> > que;
      que.emplace(ch, centroid, 1);
      while (!que.empty()) {
        int idx, par, dep;
        tie(idx, par, dep) = que.front();
        que.pop();
        if (cnt.size() <= dep) cnt.resize(dep + 1, 0);
        if (num.size() <= dep) num.resize(dep + 1, 0);
        cnt[dep]++;
        num[dep]++;
        for (auto& to : g.g[idx]) {
          if (to == par || used[to]) continue;
          que.emplace(to.to, idx, dep + 1);
        }
      }
      auto ret = FastFourierTransform::multiply(num, num);
      for (int i = 0; i < ret.size(); i++) dist[i] -= ret[i];
    }
    auto ret = FastFourierTransform::multiply(cnt, cnt);
    for (int i = 0; i < ret.size(); i++) dist[i] += ret[i];
    for (auto& to : g.tree.g[centroid]) rec(to);
  };
  rec(root);
  dist.erase(begin(dist));
  for (auto& p : dist) p /= 2;
  for (int i = 0; i < N - 1; ++i) {
    if (i) cout << ' ';
    cout << dist[i];
  }
  cout << '\n';
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	almost_uni_00	AC	212 ms	31 MB
g++	almost_uni_01	AC	214 ms	31 MB
g++	example_00	AC	2 ms	4 MB
g++	line_00	AC	767 ms	51 MB
g++	line_star_00	AC	665 ms	47 MB
g++	line_star_01	AC	544 ms	46 MB
g++	line_star_02	AC	707 ms	48 MB
g++	max_random_00	AC	366 ms	32 MB
g++	max_random_01	AC	359 ms	32 MB
g++	random_00	AC	218 ms	22 MB
g++	random_01	AC	239 ms	25 MB
g++	random_02	AC	78 ms	11 MB
g++	random_03	AC	295 ms	28 MB
g++	random_04	AC	17 ms	6 MB
g++	small_random_00	AC	3 ms	4 MB
g++	small_random_01	AC	2 ms	4 MB
g++	small_random_02	AC	2 ms	4 MB
g++	small_random_03	AC	2 ms	4 MB
g++	small_random_04	AC	2 ms	4 MB
g++	uni_00	AC	191 ms	31 MB
clang++	almost_uni_00	AC	199 ms	32 MB
clang++	almost_uni_01	AC	221 ms	32 MB
clang++	example_00	AC	2 ms	4 MB
clang++	line_00	AC	778 ms	53 MB
clang++	line_star_00	AC	668 ms	48 MB
clang++	line_star_01	AC	499 ms	47 MB
clang++	line_star_02	AC	718 ms	49 MB
clang++	max_random_00	AC	371 ms	32 MB
clang++	max_random_01	AC	360 ms	32 MB
clang++	random_00	AC	197 ms	22 MB
clang++	random_01	AC	262 ms	25 MB
clang++	random_02	AC	77 ms	11 MB
clang++	random_03	AC	306 ms	28 MB
clang++	random_04	AC	18 ms	6 MB
clang++	small_random_00	AC	3 ms	4 MB
clang++	small_random_01	AC	2 ms	4 MB
clang++	small_random_02	AC	2 ms	4 MB
clang++	small_random_03	AC	2 ms	4 MB
clang++	small_random_04	AC	2 ms	4 MB
clang++	uni_00	AC	180 ms	31 MB

Back to top page