AIOJ

Centroid decomposition

Centroid decomposition is suitable for handling path information on large trees.

Example 1: Luogu P3806 [Template] Centroid Decomposition 1

Given a weighted tree with $n$ vertices and $m$ queries, each query gives $k$ and asks whether there is a pair of vertices at distance $k$ in the tree.

$n\le 10000,m\le 100,k\le 10000000$

First, choose an arbitrary node as the root $\mathit{rt}$ . Every path entirely contained in its subtree can be divided into two types: paths passing through the current root and paths not passing through the current root. Paths passing through the current root can further be divided into paths with the root as one endpoint and paths whose two endpoints are not the root. The latter can be formed by joining two chains of the former type. Therefore, for each enumerated root $rt$ , we first calculate the contribution to the answer from paths in its subtree that pass through the node, and then recursively solve the subtrees for paths that do not pass through the node.

In this problem, for paths through the root $\mathit{rt}$ , we first enumerate all its children $\mathit{ch}$ . Taking $\mathit{ch}$ as the root, we calculate the distances from all nodes in the subtree of $\mathit{ch}$ to $\mathit{rt}$ . Let the distance from node $i$ to the current root $rt$ be $\mathit{dist}_i$ ; $\mathit{tf}_{d}$ indicates whether a previously processed subtree contains a node $v$ with $\mathit{dist}_v=d$ . If a query's $k$ satisfies $tf_{k-\mathit{dist}_i}=true$ , then a path of length $k$ exists. After determining whether the edges connected to the subtree of $\mathit{ch}$ can form an answer, we add the new distances to the $\mathit{tf}$ array.

When clearing the $\mathit{tf}$ array, do not use memset directly. Instead, put the previously occupied positions of $\mathit{tf}$ into a queue and clear them; this is necessary to maintain the time complexity.

During centroid decomposition, each point is processed once in total over all recursive calls at a given level. If there are $h$ levels, the total time complexity is $O(hn)$ .

If we choose the centroid of each subtree as its root, the number of recursive levels is minimized and the time complexity is $O(n\log n)$ . Therefore, centroid decomposition is also commonly called centroid decomposition in international competitive programming circles.

Be careful to recalculate the subtree sizes after choosing a new root; otherwise, a seemingly minor change may lead to an incorrect time complexity or make correctness difficult to guarantee.

Reference code

/* size computes subtree sizes d[], and the component size cnt
   dp computes the largest subtree f[] and the tree centroid rot
   get computes distances from nodes to the centroid t[], top
   calc performs centroid decomposition using bu[] as a length bucket
   hd, to, nx, wg store the graph in forward-star form
   ak[] and as[] handle queries offline
   ok[] marks points that have already become centroids
 */
#include <iostream>
const int N = 1e4 + 4, M = 105, Q = 1e7 + 7;
int n, m, hd[N], to[N * 2], nx[N * 2], wg[N * 2];
int ak[M], d[N], f[N], t[N], top, cnt, rot;
bool as[M], ok[N], bu[Q];

int size(int u, int pa) {
  cnt++, d[u] = 1;
  for (int p = hd[u]; ~p; p = nx[p])
    if (to[p] != pa && !ok[to[p]]) d[u] += size(to[p], u);
  return d[u];
}

void dp(int u, int pa) {
  f[u] = cnt - d[u];
  for (int p = hd[u]; ~p; p = nx[p])
    if (to[p] != pa && !ok[to[p]]) {
      f[u] = std::max(f[u], d[to[p]]);
      dp(to[p], u);
    }
  if (f[u] < f[rot]) rot = u;
}

void get(int u, int pa, int dis) {
  t[top++] = dis;
  for (int p = hd[u]; ~p; p = nx[p])
    if (to[p] != pa && !ok[to[p]]) get(to[p], u, dis + wg[p]);
}

void calc(int u) {
  cnt = 0, size(u, u);
  rot = u, dp(u, u);
  bu[0] = true, t[0] = 0, top = 1;
  for (int p = hd[rot], i; ~p; p = nx[p])
    if (!ok[to[p]]) {
      i = top, get(to[p], rot, wg[p]);
      for (int q = 0; q < m; q++)
        for (int j = i; j < top && !as[q]; j++)
          if (ak[q] >= t[j]) as[q] = bu[ak[q] - t[j]];
      --i;
      while (++i < top)
        if (t[i] < Q) bu[t[i]] = true;
    }
  while (top--)
    if (t[top] < Q) bu[t[top]] = false;
  ok[rot] = true;
  for (int p = hd[rot]; ~p; p = nx[p])
    if (!ok[to[p]]) calc(to[p]);
}

int main() {
  std::cin >> n >> m;
  for (int i = 1; i <= n; i++) hd[i] = -1;
  for (int i = 0, u, v; i + 2 < n * 2;) {
    std::cin >> u >> v >> wg[i];
    wg[i + 1] = wg[i];
    to[i] = v, nx[i] = hd[u], hd[u] = i++;
    to[i] = u, nx[i] = hd[v], hd[v] = i++;
  }
  for (int i = 0; i < m; i++) std::cin >> ak[i];
  calc(1);
  for (int i = 0; i < m; i++) std::cout << (as[i] ? "AYE\n" : "NAY\n");
}

Example 2: Luogu P4178 Tree

Given a weighted tree with $n$ vertices and a value $k$ , find the number of pairs of vertices whose distance in the tree is at most $k$ .

$n\le 40000,k\le 20000,w_i\le 1000$

Since this asks for the number of pairs whose tree distance lies in $[0,k]$ , we use a segment tree to support updates and queries.

Reference code

#include <algorithm>
#include <cstring>
#include <iostream>
#include <queue>
using namespace std;
constexpr long long MAXN = 2000010;
constexpr long long inf = 2e9;
long long n, a, b, c, q, rt, siz[MAXN], maxx[MAXN], dist[MAXN];
long long cur, h[MAXN], nxt[MAXN], p[MAXN], w[MAXN], ret;
bool vis[MAXN];

void add_edge(long long x, long long y, long long z) {
  cur++;
  nxt[cur] = h[x];
  h[x] = cur;
  p[cur] = y;
  w[cur] = z;
}

long long sum;

void calcsiz(long long x, long long fa) {
  siz[x] = 1;
  maxx[x] = 0;
  for (long long j = h[x]; j; j = nxt[j])
    if (p[j] != fa && !vis[p[j]]) {
      calcsiz(p[j], x);
      maxx[x] = max(maxx[x], siz[p[j]]);
      siz[x] += siz[p[j]];
    }
  maxx[x] = max(maxx[x], sum - siz[x]);
  if (maxx[x] < maxx[rt]) rt = x;
}

long long dd[MAXN], cnt;

void calcdist(long long x, long long fa) {
  dd[++cnt] = dist[x];
  for (long long j = h[x]; j; j = nxt[j])
    if (p[j] != fa && !vis[p[j]])
      dist[p[j]] = dist[x] + w[j], calcdist(p[j], x);
}

queue<long long> tag;

struct segtree {
  long long cnt, rt, lc[MAXN], rc[MAXN], sum[MAXN];

  void clear() {
    while (!tag.empty()) update(rt, 1, 20000000, tag.front(), -1), tag.pop();
    cnt = 0;
  }

  void print(long long o, long long l, long long r) {
    if (!o || !sum[o]) return;
    if (l == r) {
      cout << l << ' ' << sum[o] << '\n';
      return;
    }
    long long mid = (l + r) >> 1;
    print(lc[o], l, mid);
    print(rc[o], mid + 1, r);
  }

  void update(long long& o, long long l, long long r, long long x,
              long long v) {
    if (!o) o = ++cnt;
    if (l == r) {
      sum[o] += v;
      if (!sum[o]) o = 0;
      return;
    }
    long long mid = (l + r) >> 1;
    if (x <= mid)
      update(lc[o], l, mid, x, v);
    else
      update(rc[o], mid + 1, r, x, v);
    sum[o] = sum[lc[o]] + sum[rc[o]];
    if (!sum[o]) o = 0;
  }

  long long query(long long o, long long l, long long r, long long ql,
                  long long qr) {
    if (!o) return 0;
    if (r < ql || l > qr) return 0;
    if (ql <= l && r <= qr) return sum[o];
    long long mid = (l + r) >> 1;
    return query(lc[o], l, mid, ql, qr) + query(rc[o], mid + 1, r, ql, qr);
  }
} st;

void dfz(long long x, long long fa) {
  // tf[0]=true;tag.push(0);
  st.update(st.rt, 1, 20000000, 1, 1);
  tag.push(1);
  vis[x] = true;
  for (long long j = h[x]; j; j = nxt[j])
    if (p[j] != fa && !vis[p[j]]) {
      dist[p[j]] = w[j];
      calcdist(p[j], x);
      for (long long k = 1; k <= cnt; k++)
        if (q - dd[k] >= 0)
          ret += st.query(st.rt, 1, 20000000, max(0ll, 1 - dd[k]) + 1,
                          max(0ll, q - dd[k]) + 1);
      for (long long k = 1; k <= cnt; k++)
        st.update(st.rt, 1, 20000000, dd[k] + 1, 1), tag.push(dd[k] + 1);
      cnt = 0;
    }
  st.clear();
  for (long long j = h[x]; j; j = nxt[j])
    if (p[j] != fa && !vis[p[j]]) {
      sum = siz[p[j]];
      rt = 0;
      maxx[rt] = inf;
      calcsiz(p[j], x);
      calcsiz(rt, -1);
      dfz(rt, x);
    }
}

signed main() {
  cin.tie(nullptr)->sync_with_stdio(false);
  cin >> n;
  for (long long i = 1; i < n; i++)
    cin >> a >> b >> c, add_edge(a, b, c), add_edge(b, a, c);
  cin >> q;
  rt = 0;
  maxx[rt] = inf;
  sum = n;
  calcsiz(1, -1);
  calcsiz(rt, -1);
  dfz(rt, -1);
  cout << ret << '\n';
  return 0;
}

Example 3: Luogu P2664 Tree Game

A tree is given with a color assigned to every node. Define $s(i,j)$ as the number of colors on the path from $\mathit{i}$ to $\mathit{j}$ , and let $\mathit{sum_{i}}=\sum_{j=1}^n s(i,j)$ . For every $1\leq i\leq n$ , find $sum_i$ .（ $1 \le n, c_i \le 10^5$ ）

This problem strongly tests one's understanding and application of centroid decomposition. It is suitable as a relatively difficult example and exercise for centroid decomposition.

First, we need to understand a transformation. The problem defines $\mathit{sum_i}$ as the sum of the number of colors on the paths from $i$ to all nodes. However, this is difficult to count in centroid decomposition because it is hard to merge the information from two subtrees rooted at the current root. So we transform the meaning of $\mathit{sum_i}$ . For each color $j$ , consider paths with one endpoint at $i$ that contain color $j$ ; let their count be $\mathit{cnt_j}$ . Then $\mathit{sum_i}$ is actually $\sum \mathit{cnt_j}$ . This transformation merely changes the object being observed: we consider the contribution of each color to $\mathit{sum_i}$ . The value of $\mathit{cnt_j}$ is easy to obtain; whenever we encounter a new color, simply do $\mathit{cnt_{col_u}}+=\mathit{size_u}$ , where $\mathit{size_u}$ is the size of the subtree rooted at $u$ , so every node in it contributes to the answer for this node through this color.

During centroid decomposition, we only need to count the following separately:

The contribution to the root from paths in its subtree that have the current root as an endpoint.
The contribution to every point in the subtree from paths whose lca is the current root.

Part 1 is relatively straightforward. Since the number of levels in centroid decomposition is no more than $\log{n}$ , we can traverse all subtrees at every level and use the definition of $\mathit{sum_i}$ to count the result while traversing the subtrees.

For Part 2, let the current root node be $u$ and one of its children be $d$ . Let $d$ 's subtree contain an arbitrary point $v$ ; then the answer for $v$ can be divided into two parts:

The colors appearing on the path $(u, v)$ , whose number is $\mathit{num}$ , and the total size of all subtrees of $u$ other than $d$ , denoted by $\mathit{siz1}$ . The contribution of these colors to the answer for $v$ is $\mathit{num}\times \mathit{siz1}$ .
On the path $(u, v)$ , consider a color $j$ that does not appear; its contribution comes from the subtrees of $u$ other than $d$ through $\mathit{cnt_j}$ . This part of the answer is $\sum_{j \notin (u, v)} \mathit{cnt_j}$ .

This is the complete counting strategy; see the reference code for implementation details.

Reference code

#include <algorithm>
#include <iostream>
using namespace std;
#define rep(i, a, b) for (int i = (a); i <= (b); ++i)
constexpr int N = 200005;
int h[N], nxt[N * 2], to[N * 2], c[N], gr;

void tu(int x, int y) { to[++gr] = y, nxt[gr] = h[x], h[x] = gr; }

using ll = long long;
int n, nn, siz[N], mn, rt;
bool vis[N];

void get_root(int u, int f) {
  siz[u] = 1;
  int mx = 0;
  for (int i = h[u]; i; i = nxt[i]) {
    int d = to[i];
    if (vis[d] || d == f) continue;
    get_root(d, u);
    siz[u] += siz[d];
    mx = max(mx, siz[d]);
  }
  mx = max(mx, nn - siz[u]);
  if (mx < mn) mn = mx, rt = u;
}

ll ans[N], sum;
int cnt[N], v[N];
// sum keeps the running total of cnt[i]
int nowrt;

void get_dis(int u, int f, int now) {  // now is the number of colors on the current tree chain (excluding u)
  siz[u] = 1;
  if (!v[c[u]]) {
    sum -= cnt[c[u]];  // subtract color information already present in previous subtrees
    now++;
  }
  v[c[u]]++;
  ans[u] += sum + now * siz[nowrt];  // count the contribution of paths through u to u
  for (int i = h[u]; i; i = nxt[i]) {
    int d = to[i];
    if (d == f || vis[d]) continue;
    get_dis(d, u, now);
    siz[u] += siz[d];
  }
  v[c[u]]--;
  if (!v[c[u]]) {
    sum += cnt[c[u]];  // backtrack
  }
}

void get_cnt(int u, int f) {
  if (!v[c[u]]) {
    cnt[c[u]] += siz[u];
    sum += siz[u];  // merge the information from the subtree just traversed into cnt[i] and sum
  }
  v[c[u]]++;
  for (int i = h[u]; i; i = nxt[i]) {
    int d = to[i];
    if (vis[d] || d == f) continue;
    get_cnt(d, u);
  }
  v[c[u]]--;
}

void clear(int u, int f, int now) {
  if (!v[c[u]]) now++;
  v[c[u]]++;
  ans[u] -= now;
  ans[nowrt] += now;
  for (int i = h[u]; i; i = nxt[i]) {
    int d = to[i];
    if (vis[d] || d == f) continue;
    clear(d, u, now);
  }
  v[c[u]]--;
  cnt[c[u]] = 0;
}

void clear2(int u, int f) {
  cnt[c[u]] = 0;
  for (int i = h[u]; i; i = nxt[i]) {
    int d = to[i];
    if (vis[d] || d == f) continue;
    clear2(d, u);
  }
}

int son[N];

void divid(int u) {
  vis[u] = true;
  int tot = 0;
  nowrt = u;
  ans[u]++;
  for (int i = h[u]; i; i = nxt[i]) {
    if (vis[to[i]]) continue;
    son[++tot] = to[i];
  }
  siz[u] = sum = cnt[c[u]] = 1;
  v[c[u]]++;
  rep(i, 1, tot) {  // count the contribution from combinations of each subtree with all preceding subtrees
    int d = son[i];
    get_dis(d, u, 0);
    get_cnt(d, u);
    siz[u] += siz[d];
    cnt[c[u]] += siz[d];
    sum += siz[d];
  }
  clear2(u, 0);  // clear the array; remember that memset cannot be used
  siz[u] = sum = cnt[c[u]] = 1;
  for (int i = tot; i >= 1;
       --i) {  // count the contribution from combinations of each subtree with all following subtrees
    int d = son[i];
    get_dis(d, u, 0);
    get_cnt(d, u);
    siz[u] += siz[d];
    cnt[c[u]] += siz[d];
    sum += siz[d];
  }
  v[c[u]]--;
  clear(u, 0, 0);                      // count the answer while clearing
  for (int i = h[u]; i; i = nxt[i]) {  // continue centroid decomposition downward
    int d = to[i];
    if (vis[d]) continue;
    nn = siz[d], mn = n + 1, rt = 0;
    get_root(d, u);
    divid(rt);
  }
}

int main() {
  cin.tie(nullptr)->sync_with_stdio(false);
  cin >> n;
  int u, v;
  rep(i, 1, n) cin >> c[i];
  rep(i, 2, n) cin >> u >> v, tu(u, v), tu(v, u);
  rt = 0, nn = n, mn = n + 1;
  get_root(1, 0);
  divid(rt);
  rep(i, 1, n) cout << ans[i] << '\n';
  return 0;
}

Edge divide and conquer

Similar to the centroid decomposition above, we choose an edge and divide the tree as evenly as possible into two parts, making the $\mathit{size}$ values of the two subtrees connected by the edge as close as possible. We then recursively process the left and right subtrees and collect the information.

However, this does not work. Consider a flower graph:

Flower graph

When a node has multiple children with similar $\mathit{size}$ values, we find that the time complexity of applying edge divide and conquer is unacceptable.

If the graph is a binary tree, the drawback of applying edge divide and conquer to the flower graph above can be avoided. Therefore, we consider converting a multiway tree into a binary tree.

Clearly, we only need to build the tree as we would build a segment tree. It looks like this:

Building the tree

Assign suitable information to the newly created nodes according to the problem. For example, when counting path lengths, assign weight $1$ to each original edge and weight $0$ to each newly created edge.

Analyzing the complexity, we find that at most $O(n)$ nodes are added, so the total complexity is $O(n\log n)$ .

Almost all problems that can be solved by centroid decomposition can also be solved by edge divide and conquer (the constants differ, but they are not usually the bottleneck), so no example is provided here.

Centroid tree

A centroid tree is a reconstructed tree obtained by changing the shape of the original tree so that its number of levels becomes a stable $\log n$ .

It is commonly used to solve problems with modifications whose solutions do not depend on the original shape of the tree.

Algorithm analysis

We reconstruct the original tree by finding a centroid through centroid decomposition at each step.

Connect each centroid found to the centroid in the previous level as parent and child, forming a tree with $\log n$ levels.

Because the tree has $\log n$ levels, many brute-force approaches that would not have an appropriate complexity on the original tree have the correct complexity on the centroid tree.

Code implementation

There is a small trick: each time, subtract the size of the heavy child of the node in the previous level from the total size $\mathit{tot}$ of the previous level; the result is the total size of the current level. This means that finding the centroid requires only one DFS.

Reference code

#include <algorithm>
#include <iostream>
#include <vector>
using namespace std;

using IT = vector<int>::iterator;

struct Edge {
  int to, nxt, val;

  Edge() {}

  Edge(int to, int nxt, int val) : to(to), nxt(nxt), val(val) {}
} e[300010];

int head[150010], cnt;

void addedge(int u, int v, int val) {
  e[++cnt] = Edge(v, head[u], val);
  head[u] = cnt;
}

int siz[150010], son[150010];
bool vis[150010];

int tot, lasttot;
int maxp, root;

void getG(int now, int fa) {
  siz[now] = 1;
  son[now] = 0;
  for (int i = head[now]; i; i = e[i].nxt) {
    int vs = e[i].to;
    if (vs == fa || vis[vs]) continue;
    getG(vs, now);
    siz[now] += siz[vs];
    son[now] = max(son[now], siz[vs]);
  }
  son[now] = max(son[now], tot - siz[now]);
  if (son[now] < maxp) {
    maxp = son[now];
    root = now;
  }
}

struct Node {
  int fa;
  vector<int> anc;
  vector<int> child;
} nd[150010];

int build(int now, int ntot) {
  tot = ntot;
  maxp = 0x7f7f7f7f;
  getG(now, 0);
  int g = root;
  vis[g] = true;
  for (int i = head[g]; i; i = e[i].nxt) {
    int vs = e[i].to;
    if (vis[vs]) continue;
    int tmp = build(vs, ntot - son[vs]);
    nd[tmp].fa = now;
    nd[now].child.push_back(tmp);
  }
  return g;
}

int virtroot;

int main() {
  int n;
  cin >> n;
  for (int i = 1; i < n; i++) {
    int u, v, val;
    cin >> u >> v >> val;
    addedge(u, v, val);
    addedge(v, u, val);
  }
  virtroot = build(1, n);
}

Tree divide and conquer

Centroid decomposition

Edge divide and conquer

Centroid tree

Algorithm analysis

Code implementation