U41492 树上数颜色
大意
给一棵有 \(200000\) 点的树，第 \(u\) 个节点，染了颜色 \(c[u]\)，然后问
- 每个点子树内，颜色种类数。
- 第 \(u\) 个节点子树内颜色为 \(x\) 的点有多少个。
思路
经典的启发式合并题目。
由于我们总是把小的合并到大的，对于节点 \(i\)，每当它被移动一次（从集合 \(A\) 搬到 \(B\)），它所属的集合大小 \(S\) 必定满足：
\[S_{new} = |A| + |B|
\]
因为是启发式合并，必有 \(|B| \ge |A|\)，所以：
\[S_{new} \ge 2 \cdot |A| \ge 2 \cdot S_{old}
\]
设节点 \(i\) 被移动了 \(k_i\) 次。
第 1 次移动后，所在集合大小 \(\ge 2^1\)
第 2 次移动后，所在集合大小 \(\ge 2^2\)
...
第 \(k_i\) 次移动后，所在集合大小 \(\ge 2^{k_i}\)
因为集合大小上限是全树节点总数 \(n\)：
\[2^{k_i} \le n
\]
两边取对数：
\[k_i \le \log_2 n
\]
总代价是所有点移动次数之和：
\[\sum_{i=1}^{n} k_i \le \sum_{i=1}^{n} \log_2 n = n \log_2 n
\]
代码
#include <cstdio>
#include <map>
#include <vector>
using namespace std;
const int N = 100000 + 10;
int n, c[N];
vector<int> g[N];
int ans[N];
map<int, int> mp[N];void dfs(int u, int p){mp[u][c[u]] = 1;for(int i = 0;i < g[u].size();i ++){int v = g[u][i];if(v == p) continue;dfs(v, u);if(mp[v].size() > mp[u].size()){swap(mp[u], mp[v]);}for(map<int, int>::iterator it = mp[v].begin();it != mp[v].end();it ++){mp[u][it -> first] += it -> second;}mp[v].clear();}ans[u] = mp[u].size();
}int main() { scanf("%d", &n);for(int i = 2;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", &c[i]);}dfs(1, 1);for(int i = 1;i <= n;i ++){printf("%d\n", ans[i]);}return 0;
}