树上启发式合并 学习笔记

树上启发式合并 (dsu on tree) 可以在较优秀的复杂度内解决某些树上离线问题。

引入

考虑这样一个问题：

• 给一棵 $n$ 个点的有根树，求每棵子树的颜色种类数。

考虑暴力，当我们要求 $x$ 的子树的答案时，我们可以用一个 $c$ 数组记录每个颜色出现的次数，将 $x$ 子树遍历一遍后，$c$ 数组中不为 $0$ 的位置的个数就是 $x$ 子树的颜色种类数，然后清空 $c$ 数组。

显然这样总复杂度是 $\mathcal O(n^2)$ 的，有没有办法优化一下呢？

dsu on tree

当我们求完 $x$ 子树的答案后，如果我们立刻要求 $fa_x$ 子树的答案，那么我们就不必清空 $c$ 数组，同时遍历 $fa_x$ 子树时也不需要遍历 $x$ 子树。

看起来这样似乎并没有优化多少？

那么，如果我们对于每个点 $x$，在求完其重儿子 $son_x$ 的答案后，立刻按照上述方法求 $x$ 的答案，你可以惊奇的发现时间复杂度被优化到了 $\mathcal O(n \log n)$！

具体实现见代码：

const int N = 1e5 + 7;
int n, m, a[N], s[N], son[N], dfn[N], num, c[N], now, ans[N];
vi e[N];

void dfs(int x, int f) {
	s[x] = 1, dfn[x] = ++num;
	for (auto y : e[x]) if (y != f) dfs(y, x), s[x] += s[y], son[x] = s[y] > s[son[x]] ? y : son[x];
}

inline void add(int x) {
	now += !c[x]++;
}

void dfs(int x, int f, bool k) {
	for (auto y : e[x]) if (y != f && y != son[x]) dfs(y, x, 0);
	if (son[x]) dfs(son[x], x, 1);
	for (auto y : e[x]) if (y != f && y != son[x]) for (int i = 0; i < s[y]; i++) add(a[dfn[y]+i]);
	add(a[dfn[x]]), ans[x] = now;
	if (!k) for (int i = now = 0; i < s[x]; i++) c[a[dfn[x]+i]] = 0;
}

int main() {
	rd(n);
	for (int i = 1, x, y; i < n; i++) rd(x), rd(y), e[x].pb(y), e[y].pb(x);
	dfs(1, 0);
	for (int i = 1; i <= n; i++) rd(a[dfn[i]]);
	dfs(1, 0, 1), rd(m);
	for (int i = 1, x; i <= m; i++) rd(x), print(ans[x]);
	return 0;
}

证明

首先我们知道一个点到根节点上有 $\mathcal O(\log n)$ 条重链和 $\mathcal O(\log n)$ 条轻边。

注意到每条轻边会导致下面的子树被多遍历一次，因此每个点最多被遍历 $\mathcal O(\log n)$ 次。

因此总时间复杂度为 $\mathcal O(n \log n)$。

【例题】CF600E Lomsat gelral

求子树众数和，就在上面代码的基础上改一下即可。

const int N = 1e5 + 7;
int n, a[N], b[N], s[N], son[N], dfn[N], num, c[N], mx;
ll now, ans[N];
vi e[N];

void dfs(int x, int f) {
	s[x] = 1, dfn[x] = ++num;
	for (auto y : e[x]) if (y != f) dfs(y, x), s[x] += s[y], son[x] = s[y] > s[son[x]] ? y : son[x];
}

inline void add(int x) {
	if (++c[x] > mx) mx = c[x], now = 0;
	if (c[x] == mx) now += x;
}

void dfs(int x, int f, bool k) {
	for (auto y : e[x]) if (y != f && y != son[x]) dfs(y, x, 0);
	if (son[x]) dfs(son[x], x, 1);
	for (auto y : e[x]) if (y != f && y != son[x]) for (int i = 0; i < s[y]; i++) add(a[dfn[y]+i]);
	add(a[dfn[x]]), ans[x] = now;
	if (!k) for (int i = mx = now = 0; i < s[x]; i++) c[a[dfn[x]+i]] = 0;
}

int main() {
	rd(n);
	for (int i = 1; i <= n; i++) rd(b[i]);
	for (int i = 1, x, y; i < n; i++) rd(x), rd(y), e[x].pb(y), e[y].pb(x);
	dfs(1, 0);
	for (int i = 1; i <= n; i++) a[dfn[i]] = b[i];
	dfs(1, 0, 1);
	for (int i = 1; i <= n; i++) print(ans[i], " \n"[i==n]);
	return 0;
}

【练习】CF741D Arpa’s letter-marked tree and Mehrdad’s Dokhtar-kosh paths

【练习】CF570D Tree Requests

参考资料

• OI Wiki 树上启发式合并
• p-z-y dsu on tree 学习笔记