DSU on Tree

Introduction

What is a heuristic algorithm?

Heuristic algorithms are optimizations based on human experience and intuition for certain algorithms.

For example, the most common one is the heuristic merge in Disjoint Set Union (DSU). The code looks like this:

void merge(int x, int y) {
  int xx = find(x), yy = find(y);
  if (size[xx] < size[yy]) swap(xx, yy);
  fa[yy] = xx;
  size[xx] += size[yy];
}

Here, for two sets of different sizes, we merge the smaller set into the larger one, rather than merging the larger set into the smaller one.

Why? The size of a set can be considered as its height (under normal circumstances), and merging the set with smaller height into the one with larger height obviously helps us find the parent.

Making the tree with smaller height a subtree of the tree with larger height, this optimization can be called the heuristic merge algorithm.

Algorithm Overview

DSU on Tree (dsu on tree) is an algorithm for certain offline tree problems that is faster than or equal to most algorithms and easier to understand and implement.

Consider the following problem: Counting Colors on a Tree.

Example Problem

Given a tree with \(n\) nodes rooted at \(1\), node \(u\) has color \(c_u\). For each node \(u\), ask how many distinct colors appear in the subtree rooted at \(u\).

\(n\le 2\times 10^5\).

For such problems, solutions often involve heavy use of data structures (like tree of trees, etc.). If offline processing is allowed, is there a simpler method?

Process

Since offline queries are supported, consider preprocessing and then outputting answers in \(O(1)\).

Direct brute-force preprocessing has a time complexity of \(O(n^2)\), i.e., traversing each subtree once. The complexity of each traversal is clearly of the same order as \(n\), and with \(n\) nodes, the total complexity is \(O(n^2)\).

We can observe that the answer for each node is derived from its subtree and itself. Consider leveraging this property to solve the problem.

We can first precompute the size of each node's subtree and its heavy son (similar to heavy-light decomposition). The heavy son is the child with the largest subtree. This process can obviously be done in \(O(n)\).

Let \(cnt_i\) denote the occurrence count of color \(i\), and \(ans_u\) denote the answer for node \(u\).

To traverse a node \(u\), we follow these steps:

1. First traverse the light (non-heavy) sons of \(u\), compute their answers, but do not retain their impact on the \(cnt\) array;
2. Traverse its heavy son, retain its impact on the \(cnt\) array;
3. Traverse the light sons' subtrees again, adding their contributions to obtain the answer for \(u\).

The above figure is an example.

Thus, for a node, we traverse its heavy subtree once and its light subtrees twice, which is clearly optimal.

By executing this process, we obtain answers for all subtrees of this node.

Why not merge steps 1 and 3? Because the \(cnt\) array cannot be reused; otherwise, the space would be too large. We need to complete the process within \(O(n)\) space.

Clearly, if a node \(u\) is traversed \(x\) times, its heavy son will be traversed \(x\) times, and its light sons (if any) will be traversed \(2x\) times.

Note: except for the heavy son, the \(cnt\) array must be cleared after each traversal.

Proof

We define heavy edges and light edges as in heavy-light decomposition (edges connecting to heavy sons are heavy edges, others are light edges). For the definitions of heavy son and heavy edge, see the figure below. For a tree with \(n\) nodes:

The number of light edges from the root to any node on the tree does not exceed \(\log n\). Suppose there are \(x\) light edges from the root to a node, and the subtree size of that node is \(y\). Clearly, the subtree size of a child connected by a light edge is less than half of its parent's subtree size (otherwise it wouldn't be a light edge). Thus, \(y < n/2^x\), and obviously \(n > 2^x\), so \(x < \log n\).

Moreover, if a node is the heavy son of its parent, its subtree must be the largest among its siblings. Therefore, any node on the path from an arbitrary node to the root, when computing answers, will not traverse this node via heavy edges. So the number of times a node is traversed equals the number of light edges on the path from it to the root plus \(1\) (the \(+1\) is because the node itself is traversed). Hence, the number of traversals for a node \(= \log n + 1\), and the total time complexity is \(O(n(\log n+1)) = O(n\log n)\). Outputting answers costs \(O(m)\).

The bold edges in the figure are heavy edges, and the child nodes connected by heavy edges are heavy sons.

Optimization

As mentioned in the proof, dsu on tree utilizes the concept of light and heavy sons from heavy-light decomposition to accelerate merging. Therefore, we can also directly use the dfs order obtained from heavy-light decomposition, transforming recursion into iteration, further optimizing the constant factor of dsu on tree.

The dfs order itself has the following property: the subtree of a node is always contiguous in the dfs order. Thus, we can traverse the dfs order array in reverse. This ensures that when processing a node, all other nodes in its subtree have already been processed.

The dfs order obtained from heavy-light decomposition has the following excellent property: a heavy chain is always contiguous in the dfs order. Therefore, when traversing nodes in reverse dfs order, for a node at the top of a heavy chain, the next node to traverse is definitely not its parent, so we need to clear its influence. Otherwise, for nodes not at the top of a heavy chain, the previous node traversed is either its heavy son or a node from another branch whose influence has already been cleared, so we can directly inherit its influence. Based on this, we can quickly collect the influence of all light sons using the dfs order and record the answer.

The above process is called the non-recursive/iterative implementation of dsu on tree (also known as the dfs order implementation of dsu on tree). Compared to the original recursive implementation, it reduces the time and space overhead of recursive function calls, achieving significant constant optimization, especially when dealing with trees containing many chain-like structures, offering notable stack space advantages.

Implementation

Reference Implementation

Recursive ImplementationNon-Recursive Implementation

#include <cstdio>
#include <vector>
using namespace std;

constexpr int N = 2e5 + 5;

int n, m;

// g[u]: stores nodes adjacent to u
vector<int> g[N];

// sz: subtree size
// big: heavy son
// col: node color
// L[u]: DFS order of node u
// R[u]: maximum DFS order among nodes in the subtree of u
// Node[i]: node with DFS order i
// totdfn: node counter, also the current maximum DFS order of visited nodes
// ans: stores answers
// cnt[i]: count of nodes with color i
// totColor: current number of distinct colors that have appeared
int sz[N], big[N], col[N], L[N], R[N], Node[N], totdfn;
int ans[N], cnt[N], totColor;

void add(int u) {
  if (cnt[col[u]] == 0) ++totColor;
  cnt[col[u]]++;
}

void del(int u) {
  cnt[col[u]]--;
  if (cnt[col[u]] == 0) --totColor;
}

int getAns() { return totColor; }

void dfs0(int u, int p) {
  L[u] = ++totdfn;
  Node[totdfn] = u;
  sz[u] = 1;
  for (int v : g[u])
    if (v != p) {
      dfs0(v, u);
      sz[u] += sz[v];
      if (!big[u] || sz[big[u]] < sz[v]) big[u] = v;
    }
  R[u] = totdfn;
}

void dfs1(int u, int p, bool keep) {
  // Calculate answers for light sons
  for (int v : g[u])
    if (v != p && v != big[u]) {
      dfs1(v, u, false);
    }
  // Calculate answer for heavy son and retain data from the computation process (for inheritance)
  if (big[u]) {
    dfs1(big[u], u, true);
  }
  for (int v : g[u])
    if (v != p && v != big[u]) {
      // DFS orders of subtree nodes form a contiguous interval, can be traversed directly
      for (int i = L[v]; i <= R[v]; i++) {
        add(Node[i]);
      }
    }
  add(u);
  ans[u] = getAns();
  if (!keep) {
    for (int i = L[u]; i <= R[u]; i++) {
      del(Node[i]);
    }
  }
}

int main() {
  scanf("%d", &n);
  for (int i = 1; i < n; i++) {
    int u, v;
    scanf("%d%d", &u, &v);
    g[u].push_back(v);
    g[v].push_back(u);
  }
  for (int i = 1; i <= n; i++) scanf("%d", &col[i]);
  dfs0(1, 0);
  dfs1(1, 0, false);
  scanf("%d", &m);
  for (int i = 1; i <= m; i++) {
    int q;
    scanf("%d", &q);
    printf("%d\n", ans[q]);
  }
  return 0;
}

#include <cstdio>
#include <vector>
using namespace std;

constexpr int N = 2e5 + 5;

int n, m;

// g[u]: stores nodes adjacent to u
vector<int> g[N];

// sz: subtree size
// son[u]: the heavy son of node u
// col: node color
// dfn[u]: DFS order of node u
// bottom[u]: the maximum DFS order among nodes in the subtree of u
// totdfn: node counter, also the current maximum DFS order of visited nodes
// fa[u]: parent node of u
// top[u]: the top node of the heavy chain containing node u
// rnk[i]: the node with DFS order i
// ans[u]: stores the answer
// cnt[i]: count of nodes with color i
// totColor: current number of distinct colors that have appeared
int sz[N], son[N], col[N], dfn[N], fa[N], bottom[N], totdfn;
int top[N], rnk[N];
int ans[N], cnt[N], totColor;

void add(int u) {
  if (cnt[col[u]] == 0) ++totColor;
  cnt[col[u]]++;
}

void del(int u) {
  cnt[col[u]]--;
  if (cnt[col[u]] == 0) --totColor;
}

int getAns() { return totColor; }

void dfs0(int u, int p) {
  sz[u] = 1;
  fa[u] = p;
  for (int v : g[u])
    if (v != p) {
      dfs0(v, u);
      sz[u] += sz[v];
      if (!son[u] || sz[son[u]] < sz[v]) son[u] = v;
    }
}

void dfs1(int u, int t) {
  rnk[bottom[u] = dfn[u] = ++totdfn] = u;
  top[u] = t;
  if (!son[u]) return;
  dfs1(son[u], t);
  for (int v : g[u])
    if (v != fa[u] && v != son[u]) dfs1(v, v);
  bottom[u] = totdfn;
}

void dsu_on_tree() {
  for (int i = totdfn; i >= 1; i--) {
    int u = rnk[i];
    for (int v : g[u])
      if (v != fa[u] && v != son[u]) {
        for (int i = dfn[v]; i <= bottom[v]; i++) {
          add(rnk[i]);
        }
      }
    add(u);
    ans[u] = getAns();
    if (u == top[u]) {
      for (int i = dfn[u]; i <= bottom[u]; i++) {
        del(rnk[i]);
      }
    }
  }
}

int main() {
  scanf("%d", &n);
  for (int i = 1; i < n; i++) {
    int u, v;
    scanf("%d%d", &u, &v);
    g[u].push_back(v);
    g[v].push_back(u);
  }
  for (int i = 1; i <= n; i++) scanf("%d", &col[i]);
  dfs0(1, 1);
  dfs1(1, 1);
  dsu_on_tree();
  scanf("%d", &m);
  for (int i = 1; i <= m; i++) {
    int q;
    scanf("%d", &q);
    printf("%d\n", ans[q]);
  }
  return 0;
}

Applications

1. Problems where the intended solution by some problem setters is dsu on tree

   For example, CF741D. Given a tree, each node's weight is a letter from 'a' to 'v'. Each query asks to find a path within a subtree such that the characters on the path, when sorted, form a palindrome.

   Since it becomes a palindrome after sorting, a character appearing twice is equivalent to not appearing. That is, the path must satisfy at most one character appears an odd number of times.

   The normal approach is to dfs from each node, forcibly enumerating all letters at each node to find paths where the XOR result with it has more than one bit set to 1, then take the maximum length. This is \(O(n^2\log n)\), but can be optimized to \(O(n\log^2n)\) using dsu on tree. For specific details, refer to the extended reading below.

2. Problems that can be solved with dsu tricks

   Can be used to score partial points on some tree-of-trees problems (without modification operations). Moreover, the complexity of dsu is better than the \(O(n\sqrt{m})\) of Mo's algorithm on trees.

Practice Problems

CF600E Lomsat gelral

Problem translation: Tree nodes have colors. A color dominates a subtree if and only if no other color appears more frequently in that subtree. Find the sum of all colors dominating each subtree.

UOJ284 Happy Gaming Chicken

CF1709E XOR Tree

References/Further Reading

CF741D Author's introduction to dsu on tree

This author's solution

---

本页面最近更新：，更新历史
发现错误？想一起完善？ 在 GitHub 上编辑此页！
本页面贡献者：OI-wiki
本页面的全部内容在 CC BY-SA 4.0 和 SATA 协议之条款下提供，附加条款亦可能应用