樹上啟發式合併（dsu on tree）學習筆記
樹上啟發式合併（dsu on tree）學習筆記
閒話
樹上啟發式合併，又稱 dsu on tree（雖然跟 dsu 並查集完全沒關係），用於離線處理子樹相關詢問。
它是一種利用了重鏈剖分性質的暴力，時間複雜度為完全正確的 \(\mathcal{O}(n\log n+m)\)，個人認為跟莫隊等都是非常優雅的暴力。
閱讀本文並不需要重鏈剖分作為前置知識。
記號約定
本文中用 \(u\to v\) 表示 \(v\) 是 \(u\) 的兒子，\(u\leadsto v\) 表示 \(v\) 是 \(u\) 的後代。
例子：樹上數顏色
這是樹上啟發式合併最基礎的應用。
考慮暴力，每次詢問搜一遍子樹，統計答案，但這是 \(\mathcal{O}(mn)\)
就沒有更好的做法嗎？
我們設 \({buc}_u\) 為 \(u\) 子樹內的桶，\({buc}_{u,i}\) 也就是顏色 \(i\) 在 \(u\) 子樹中出現的次數，也就是 \({buc}_{u,i}=\sum\limits_{u\leadsto v}[c_v=i]\)。
可以發現 \({buc}_{u,i}=\sum\limits_{u\to v}{buc}_{v,i}+[c_u=i]\)，但我們前面的暴力中每個節點都重複統計了 \(buc\)，是不是有些浪費呢？考慮對 \(buc\) 陣列進行重複利用。
由於兩棵沒有包含關係的子樹的 \(buc\) 不能重複利用，也不可能真的給每個節點都開一個 \(buc\)，因此對於每個節點，\(buc\) 陣列只能從一個兒子處繼承。那從哪裡繼承呢？用尾椎骨想一想就知道顯然從重兒子處繼承是最優的！這裡我們把 \(u\) 的重兒子定義為，在 \(u\) 的所有兒子中，子樹大小最大的。
於是演算法流程就出來了：
- 預處理每個點的重兒子，記點 \(u\) 的重兒子為 \({son}_u\)。
- 從根遞迴下去，對於節點 \(u\)，先統計所有輕兒子（非重兒子）的答案，並從 \(buc\) 中擦去它們的貢獻。此時的 \(buc\) 為空。
- 然後統計重兒子的答案，但從 \(buc\)中不擦去貢獻。此時的 \({buc}_i=\sum\limits_{{son}_u\leadsto v}[c_v=i]\)。
- 在 \(buc\) 中添加當前點 \(u\) 和所有輕兒子子樹內的點的貢獻。此時的 \({buc}_i=\sum\limits_{u\leadsto v}[c_v=i]\)。
- 處理出當前點的答案。
- 如果當前點是父親的輕兒子，則需要擦去貢獻，列舉 \(u\leadsto v\) 並把 \({buc}_{c_v}\) 減一即可。
那麼複雜度是啥呢？
我們稱一個節點與它的重兒子連線的邊為“重邊”，與它的輕兒子連線的邊為“輕邊”。
首先丟擲一條引理：根節點到任意節點路徑上的輕邊不超過 \(\log n\) 條。
證明：設根節點到當前節點的輕邊為 \(x\) 條，當前節點的子樹大小為 \(y\)。由輕重兒子定義，顯然輕兒子的子樹大小不超過父親的一半，則有 \(y < \dfrac{n}{2^x}\)，所以 \(n > 2^x\)，即 \(x < \log n\)。證畢。
然後考慮一個點會被計算多少次，顯然只有搜到這個點，或者這個點位於它的某個祖先的輕兒子的子樹內，這個點才會被計算。又因為上面的引理，顯然每個點被計算次數為 \(\mathcal{O}(\log n)\)，於是複雜度為 \(\mathcal{O}(n\log n)\)。
回到這道樹上數顏色，給出參考程式碼：
//By: Luogu@rui_er(122461)
#include <bits/stdc++.h>
#define rep(x,y,z) for(int x=y;x<=z;x++)
#define per(x,y,z) for(int x=y;x>=z;x--)
#define debug printf("Running %s on line %d...\n",__FUNCTION__,__LINE__)
#define fileIO(s) do{freopen(s".in","r",stdin);freopen(s".out","w",stdout);}while(false)
using namespace std;
typedef long long ll;
const int N = 1e5+5;
int n, m, c[N], sz[N], son[N], buc[N], ans[N], now;
vector<int> e[N];
template<typename T> void chkmin(T& x, T y) {if(x > y) x = y;}
template<typename T> void chkmax(T& x, T y) {if(x < y) x = y;}
void dfs(int u, int f) { // 第一步：預處理重兒子
sz[u] = 1;
for(auto v : e[u]) {
if(v == f) continue;
dfs(v, u);
sz[u] += sz[v];
if(sz[v] > sz[son[u]]) son[u] = v;
}
}
void add(int u, int f, int dt) {
if(dt > 0 && !buc[c[u]]) ++now;
buc[c[u]] += dt;
if(dt < 0 && !buc[c[u]]) --now;
for(auto v : e[u]) {
if(v == f) continue;
add(v, u, dt);
}
}
void calc(int u, int f, int sv) {
for(auto v : e[u]) { // 第二步：統計輕兒子答案並擦去貢獻
if(v == f || v == son[u]) continue;
calc(v, u, 0);
}
if(son[u]) calc(son[u], u, 1); // 第三步：統計重兒子答案並保留貢獻
if(!buc[c[u]]) ++now;
++buc[c[u]];
for(auto v : e[u]) { // 第四步：新增輕兒子貢獻
if(v == f || v == son[u]) continue;
add(v, u, 1);
}
ans[u] = now; // 第五步：處理當前點答案
if(!sv) add(u, f, -1); // 第六步：如果是輕兒子，擦去貢獻
}
int main() {
scanf("%d", &n);
rep(i, 1, n-1) {
int u, v;
scanf("%d%d", &u, &v);
e[u].push_back(v);
e[v].push_back(u);
}
rep(i, 1, n) scanf("%d", &c[i]);
dfs(1, 0);
calc(1, 0, 0);
for(scanf("%d", &m);m;m--) {
int u;
scanf("%d", &u);
printf("%d\n", ans[u]);
}
return 0;
}
存幾道題
CF600E Lomsat gelral、CF208E Blood Cousins、CF741D Arpa’s letter-marked tree and Mehrdad’s Dokhtar-kosh paths。