説明

木の重心分解を行う。

計算量

  • $O(N \log N)$

実装例

  • CentroidDecomposition($g$):= グラフ $g$ で初期化する。
  • build($t$):= 重心分解して重心を返す。$t$ には重心と、その重心を取り除いたときのサブグラフの重心たち同士を結んだ木(?)が格納される。またbelongには各頂点についてどの重心のサブグラフに所属するかが格納される。
template< typename G >
struct CentroidDecomposition {
  const G &g;
  vector< int > sub;
  vector< vector< int > > belong;
  vector< bool > v;

  CentroidDecomposition(const G &g) : g(g), sub(g.size()), v(g.size()), belong(g.size()) {}

  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 void belong_dfs(int idx, int par, int centroid) {
    belong[idx].emplace_back(centroid);
    for(auto &to : g[idx]) {
      if(to == par || v[to]) continue;
      belong_dfs(to, idx, centroid);
    }
  }

  inline int build(UnWeightedGraph &t, int idx) {
    int centroid = search_centroid(idx, -1, build_dfs(idx, -1) / 2);
    v[centroid] = true;
    belong_dfs(centroid, -1, centroid);
    for(auto &to : g[centroid]) {
      if(!v[to]) t[centroid].emplace_back(build(t, to));
    }
    v[centroid] = false;
    return centroid;
  }

  inline int build(UnWeightedGraph &t) {
    t.resize(g.size());
    return build(t, 0);
  }
};

検証

AtCoder「みんなのプロコン 2018」決勝 C 木の問題

int main() {
  int N, Q;
  scanf("%d %d", &N, &Q);
  UnWeightedGraph g(N);
  for(int i = 1; i < N; i++) {
    int x, y;
    scanf("%d %d", &x, &y);
    --x, --y;
    g[x].emplace_back(y);
    g[y].emplace_back(x);
  }
  vector< pair< int, int > > ev[100000];
  for(int i = 0; i < Q; i++) {
    int v, k;
    scanf("%d %d", &v, &k);
    --v;
    ev[v].emplace_back(k, i);
  }
  vector< int > ans(Q);
  CentroidDecomposition< UnWeightedGraph > cd(g);
  UnWeightedGraph tree;
  int root = cd.build(tree);
  vector< int > used(N), sz(N);
 
  function< void(int, int, int, int) > add_path = [&](int idx, int par, int dep, int d) {
    sz[dep] += d;
    for(auto &to : g[idx]) {
      if(to == par || used[to]) continue;
      add_path(to, idx, dep + 1, d);
    }
  };
  function< void(int, int, int) > get_path = [&](int idx, int par, int dep) {
    for(auto &e : ev[idx]) {
      int rest = e.first - dep;
      if(rest < 0) continue;
      ans[e.second] += sz[rest];
    }
    for(auto &to : g[idx]) {
      if(to == par || used[to]) continue;
      get_path(to, idx, dep + 1);
    }
  };
  function< void(int) > beet = [&](int idx) {
    used[idx] = true;
    add_path(idx, -1, 0, 1);
    for(auto &e : ev[idx]) {
      int rest = e.first;
      if(rest < 0) continue;
      ans[e.second] += sz[rest];
    }
    for(auto &to : g[idx]) {
      if(used[to]) continue;
      add_path(to, idx, 1, -1);
      get_path(to, idx, 1);
      add_path(to, idx, 1, 1);
    }
    add_path(idx, -1, 0, -1);
    for(auto &to : tree[idx]) beet(to);
    used[idx] = false;
  };
  beet(root);
  for(auto &x : ans) printf("%d\n", x);
}