题意
给出一颗以 1 为根的树,每个点有颜色,如果某个子树上某个颜色出现的次数最多,则认为它在这课子树有支配地位,一颗子树上,可能有多个有支配的地位的颜色,对每颗子树分别求有支配地位的颜色的和(把颜色这个权值相加)。
分析
树上启发式合并模板题。
如果暴力去搜索,显然是 \(O(n^2)\) 的算法,可以考虑优化,当我们搜索到节点 u 时,最后去搜索 u 的子节点中子树节点数量最大的子节点(树链剖分求出重儿子),并保留这个子节点所在子树的状态(颜色数量信息),这样在更新贡献的时候可以直接跳过它了。复杂度 \(O(nlogn)\)。
code
#includeusing namespace std;typedef long long ll;const int MAXN = 2e5 + 10;int n;int fa[MAXN], son[MAXN], dep[MAXN], siz[MAXN];int col[MAXN];int cnt, head[MAXN];struct Edge { int to, next;}e[MAXN];void addedge(int u, int v) { e[cnt].to = v; e[cnt].next = head[u]; head[u] = cnt++;}void dfs(int u) { siz[u] = 1; son[u] = 0; for(int i = head[u]; ~i; i = e[i].next) { if(e[i].to != fa[u]) { fa[e[i].to] = u; dep[e[i].to] = dep[u] + 1; dfs(e[i].to); if(siz[e[i].to] > siz[son[u]]) son[u] = e[i].to; siz[u] += siz[e[i].to]; } }}int vis[MAXN];ll sum, mx, C[MAXN];ll ans[MAXN];void update(int u, int c) { C[col[u]] += c; if(c > 0 && C[col[u]] >= mx) { if(C[col[u]] > mx) { sum = 0; mx = C[col[u]]; } sum += col[u]; } for(int i = head[u]; ~i; i = e[i].next) { if(e[i].to != fa[u] && !vis[e[i].to]) update(e[i].to, c); }}void dfs1(int u, int flg) { for(int i = head[u]; ~i; i = e[i].next) { if(e[i].to != fa[u] && e[i].to != son[u]) dfs1(e[i].to, 1); } if(son[u]) { dfs1(son[u], 0); vis[son[u]] = 1; } update(u, 1); ans[u] = sum; if(son[u]) vis[son[u]] = 0; if(flg) { update(u, -1); sum = 0; mx = 0; }}int main() { memset(head, -1, sizeof head); scanf("%d", &n); for(int i = 1; i <= n; i++) { scanf("%d", &col[i]); } for(int i = 1; i < n; i++) { int x, y; scanf("%d%d", &x, &y); addedge(x, y); addedge(y, x); } dfs(1); dfs1(1, -1); for(int i = 1; i <= n; i++) { printf("%I64d%c", ans[i], i == n ? '\n' : ' '); } return 0;}