DSU on Tree

Introduction

What is a heuristic algorithm?

A heuristic algorithm optimizes an algorithm using human experience and intuition.

The most common example is heuristic merging in a disjoint-set union; its code is as follows:

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 instead of merging the larger set into the smaller one.

Why? The size of a set can be regarded as the height of the set (under normal circumstances), and making the shorter tree a subtree of the taller tree clearly helps us find the parent.

Making the shorter tree a subtree of the taller tree is an optimization called heuristic merging.

Algorithm

DSU on tree is an algorithm for some offline problems on trees that is at least as fast as most alternatives while being easier to understand and implement.

Consider the following problem: Distinct Colors on a Tree.

Problems like this are often solved with elaborate data structures such as a tree of trees. If the queries can be processed offline, is there a simpler method?

Procedure

Since offline processing is available, consider preprocessing the answers so that each answer can be output in $O(1)$ time.

The direct brute-force preprocessing takes $O(n^2)$ time: for every node, traverse once, each traversal clearly has complexity on the order of $n$, and there are $n$ nodes, so the total complexity is $O(n^2)$.

Observe that the answer for each node is determined by its subtree and by the node itself. We can exploit this property.

First preprocess the size of every subtree and the heavy child of every node. As in heavy-light decomposition, the heavy child is the child whose subtree contains the most nodes. This can clearly be done in $O(n)$ time.

Let $cnt_i$ denote the number of occurrences of color $i$, and let $ans_u$ denote the answer for node $u$.

When traversing a node $u$, follow these steps:

1. Traverse $u$'s light (non-heavy) children and compute their answers, but do not retain the effect of the traversal on the $cnt$ array;
2. Traverse its heavy child and retain its effect on the $cnt$ array;
3. Traverse the subtree nodes of $u$'s light children again, adding their contributions to obtain the answer for $u$.

The figure above shows an example.

Thus, for each node, we traverse the heavy subtree once and every non-heavy subtree twice, which is the most efficient arrangement.

By carrying out this process, we obtain the answers for all nodes in the subtree of the current node.

Why not combine the first and third steps? The $cnt$ array cannot be reused in that way, because otherwise the space usage would be too large; the algorithm must use only $O(n)$ space.

Clearly, if a node $u$ is traversed $x$ times, its heavy child is traversed $x$ times, while each light child, if present, is traversed $2x$ times.

Remember to clear $cnt$ after every traversal except for the heavy child.

Proof

As in heavy-light decomposition, define heavy and light edges (an edge leading to the heavy child is heavy, and all other edges are light). The definitions of heavy children and heavy edges are illustrated below for a tree with $n$ nodes:

The number of light edges on the path from the root to any node is at most $\log n$. Suppose the path from the root to a node contains $x$ light edges and the node's subtree has size $y$. Every light edge connects to a child whose subtree has size less than half that of its parent (otherwise it would not be a light edge), so $y<n/2^x$. Clearly $n>2^x$, hence $x<\log n$.

Moreover, if a node is its parent's heavy child, then its subtree is the largest among its siblings. Therefore, when computing answers, none of the heavy edges on the path from any node to the root causes that node to be traversed. Thus, the number of times a node is traversed equals the number of light edges on its path to the root $+1$ (the extra $+1$ accounts for traversing the node itself). Hence, the number of times a node is traversed $=\log n+1$, and the total time complexity is $O(n(\log n+1))=O(n\log n)$. Outputting the answers costs $O(m)$.

The bold edges in the figure are heavy edges; the child at the end of a heavy edge is the heavy child.

Optimization

The proof above shows that DSU on tree uses the light/heavy child concept from heavy-light decomposition to speed up merging. We can therefore directly use the DFS order produced by heavy-light decomposition, turn the recursion into iteration, and further reduce the constant factor of DSU on tree.

The DFS order itself has the following property: the nodes in a subtree form a contiguous interval in DFS order. Thus, traverse the DFS-order array in reverse. This ensures that when a node is reached, all other nodes in its subtree have already been processed.

The DFS order produced by heavy-light decomposition has another useful property: a heavy path is always contiguous in DFS order. Therefore, when traversing nodes in reverse DFS order, for a node at the top of a heavy path, the next node to traverse cannot be its parent, so its contribution must be cleared; for a node that is not at the top of a heavy path, the preceding node is either its heavy child or a node from another branch whose contribution has already been cleared, so its contribution can be inherited directly. We can then use the DFS order to quickly count the contributions of all light children and record the answer.

This process is called the non-recursive/iterative implementation of DSU on tree (also called the DFS-order implementation of DSU on tree). Compared with the original recursive implementation, it reduces the time and space overhead of recursive calls and provides a substantial constant-factor optimization, especially offering a significant stack-space advantage when processing trees containing many chain-like structures.

Implementation

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

constexpr int N = 2e5 + 5;

int n, m;

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

// sz: subtree size
// big: heavy child
// col: node color
// L[u]: DFS order of node u
// R[u]: maximum DFS order of a node in the subtree of u
// Node[i]: node with DFS order i
// totdfn: node counter and the maximum DFS order reached so far
// ans: stores answers
// cnt[i]: number of nodes with color i
// totColor: number of colors seen so far
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) {
  // Compute answers for light children
  for (int v : g[u])
    if (v != p && v != big[u]) {
      dfs1(v, u, false);
    }
  // Compute the answer for the heavy child and retain its data for inheritance
  if (big[u]) {
    dfs1(big[u], u, true);
  }
  for (int v : g[u])
    if (v != p && v != big[u]) {
      // Nodes in a subtree form a contiguous interval in DFS order, so traverse it 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]: nodes adjacent to u
vector<int> g[N];

// sz: subtree size
// son[u]: index of the heavy child of node u
// col: node color
// dfn[u]: DFS order of node u
// bottom[u]: maximum DFS order of a node in the subtree of node u
// totdfn: node counter and the maximum DFS order reached so far
// fa[u]: index of the parent of node u
// top[u]: index of the top node of the heavy path containing u
// rnk[i]: node with DFS order i
// ans[u]: stores answers
// cnt[i]: number of nodes with color i
// totColor: number of colors seen so far
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. Some problem setters make DSU on tree the intended solution.

   For example, CF741D. Given a tree in which each node is labeled with a letter from 'a' to 'v', each query asks for a path in a subtree whose characters can be rearranged into a palindrome.

   Since a character occurring twice has no effect after rearrangement into a palindrome, the path must have at most one character with an odd frequency.

   The straightforward approach runs a DFS from every node and, at each node, enumerates all letters to find paths whose XOR result is 1 and whose count exceeds 1, then takes the longest one. This takes $O(n^2\log n)$ time, but DSU on tree reduces it to $O(n\log^2n)$. See the further reading below for the details.

2. Problems that can be solved by applying DSU in a rough way.

   It can obtain partial scores on some tree-of-trees problems (when there are no modification operations), and its complexity is better than the $O(n\sqrt{m})$ complexity of Mo's algorithm on trees.

Practice problems

CF600E Lomsat gelral

Problem statement: The nodes of a tree have colors. A color occupies a subtree if and only if no other color appears more frequently in that subtree. Find the sum of all colors that occupy each subtree.

UOJ284 Happy Game Chicken

CF1709E XOR Tree

References / further reading

The author-introduced DSU on tree for CF741D

The author's solution