Original post : [Tutorial] Sack (dsu on tree)

In this post, I am trying to explain every line of dsu on tree code.

The reason is the same as blog 51139 when I first learn it.

Please tell me if something goes wrong. Sorry for bad English.

How many vertices in the subtree of vertex v has some property in $$$O(n \ log \ n)$$$ time (for all of the queries)?

For example:

Given a tree, every vertex has a color. The query is how many vertices in the subtree of vertex v are colored with color c?

First, calculate the size of the subtree of every vertex, it can be done with a single dfs $$$(sz[v] = v + all(sz[child \ of \ v]))$$$

void dfs_size(int v, int p) {
	sz[v] = 1;
	for (auto u : adj[v]) {
		if (u != p) {
			dfs_size(u, v);
			sz[v] += sz[u];
		}
	}
}

Code:

void dfs(int v, int p, bool keep) {
	int Max = -1, bigchild = -1;
	for (auto u : adj[v]) {
		if (u != p && Max < sz[u]) {
			Max = sz[u];
			bigchild = u;
		}
	}
	for (auto u : adj[v]) {
		if (u != p && u != bigchild) {
			dfs(u, v, 0);
		}
	}
	if (bigchild != -1) {
		dfs(bigchild, v, 1);
		swap(vec[v], vec[bigchild]);
	}
	vec[v].push_back(v);
	cnt[color[v]]++;
	for (auto u : adj[v]) {
		if (u != p && u != bigchild) {
			for (auto x : vec[u]) {
				cnt[color[x]]++;
				vec[v].push_back(x);
			}
		}
	}
	// there are cnt[c] vertex in subtree v color with c
	if (keep == 0) {
		for (auto u : vec[v]) {
			cnt[color[u]]--;
		}
	}
}

Explanation :

$$$cnt[c]$$$ in $$$dfs(v)$$$ will be answered for query in vertex v with the color c

each time we do dfs in the subtree of v, we will keep all nodes in the subtree of v in $$$vec[v]$$$,

update information (color count for each node) from subtree child while adding a node to $$$vec[v]$$$.

boolean keep denote if we are working on the subtree of $$$bigchild \ (keep = 1)$$$ or $$$smallchild \ (keep = 0)$$$

void dfs(int v, int p, bool keep)

let say we are at subtree v

int Max = -1, bigchild = -1;
for (auto u : adj[v]) {
	if (u != p && Max < sz[u]) {
		Max = sz[u];
		bigchild = u;
	}
}

First, we will find a big child using pre-calc $$$sz[u]$$$.

Note that the big child of vertex v is the child with max subtree size

for (auto u : adj[v]) {
	if (u != p && u != bigchild) {
		dfs(u, v, 0);
	}
}

if (bigchild != -1) {
	dfs(bigchild, v, 1);
	swap(vec[v], vec[bigchild]);
}

then call recur for small child $$$(keep = 0)$$$ then big child $$$(keep = 1)$$$

then : $$$vec[v] = vec[big \ child] + v + all(vec[small \ child])$$$

obviously, $$$vec[big \ child]$$$ and $$$vec[small \ child]$$$ are calculated before

because we only care about current subtree, we can swap $$$vector[big \ child]$$$ and $$$vector[v]$$$

so that $$$vector[v]$$$ have bigchild information. The complexity of swap vector is $$$O(1)$$$

and current cnt array contain only bigchild information (read $$$if(keep == 0)$$$ explain part below)

$$$vec[v] = vec[big \ child]$$$

vec[v].push_back(v);
cnt[color[v]]++;

next, add information of vertex v.

$$$vec[v] = vec[big \ child] + v$$$

for (auto u : adj[v]) {
	if (u != p && u != bigchild) {
		for (auto x : vec[u]) {
			cnt[color[x]]++;
			vec[v].push_back(x);
		}
	}
}

$$$vec[v] = vec[big \ child] + v + all(vec[small \ child])$$$

now that $$$vec[v]$$$ is full fill, add small child to $$$vec[v]$$$ and update cnt array also

(keep in mind that bigchild information didn't get deleted so that we only need to add small child)

if (keep == 0) {
	for (auto u : vec[v]) {
		cnt[color[u]]--;
	}
}

go down to small child, we have to delete its information about color count so that

the query answer of v's first child subtree won't affect second and so on. (because we are using global array)

That is also the reason why we have to call small child before big child recur

(to make bigchild and small child answer not overlap)

Note:

this answer doubt in the comment, feel free to ask

cnt is a global array

we process small child first clear them from the global array, solve big child then add all small child again

Because some problems like count distinct values in subtree can't be solved if we have a contribution

from other nodes when we are at some node only it's subtree children data must be in the global array

we need to have separate information for big and small children.

we can't just work in a single entity because of the time complexity otherwise it will be $$$O(n^2)$$$, see prove below

Complexity?

for (auto u : adj[v]) {
	if (u != p && u != bigchild) {
		for (auto x : vec[u]) {
			cnt[color[x]]++;
			vec[v].push_back(x);
		}
	}
}

add small child took $$$O(n \ log \ n)$$$ overall. why?

let call $$$y = vec[big \ child].size, \ x = vec[small \ child].size \ (y \geq x)$$$

let consider a vertex u $$$(1 \leq u \leq n)$$$

merge it to the bigger child so the size of vec[] of the subtree contain u

becomes $$$x + y \geq x + x = 2x$$$. So each time we merge to parent subtree

the size increase at least twice when adding and total size is only n so

we cannot add a node more than $$$log(n)$$$ time

we have n vertex u so the complexity becomes $$$O(n \ log \ n)$$$

if (keep == 0) {
	for (auto u : vec[v]) {
		cnt[color[u]]--;
	}
}

because we only deleting information of subtree when $$$keep = 0$$$ (small subtree)

which is the same as when we add it (now delete instead of add) so the complexity is $$$O(n \ log \ n)$$$

so overall complexity is $$$O(n \ log \ n)$$$

Example:

600E - Lomsat gelral

Practice:

visit original blog :>