Lomsat gelral【Dsu on Tree】

Lomsat gelral【Dsu on Tree】

  题意：有N个结点构成的一棵树，其中每个结点都有一个颜色，现在我们想知道每个结点的统治色是哪种或者哪几种颜色，如果存在多种统治色的话，就是这几种统治色的颜色权值和。

  思路：这是一道Dsu on Tree的模板题，但是作为初学Dsu on Tree的小白，我还是想分享一下我的认知。

  首先，Dsu的主心骨就是重链剖分。我们利用重链剖分将结点的点分为了轻儿子和重儿子，这样，我们的做法就是对轻儿子直接暴力，对重儿子直接O(1)的记录下来，但是为了重儿子不影响答案的正确性，我们是必须要访问所有的轻儿子之后，才能访问对应的结点的重儿子。

  那么，我们从一个结点向下，我们肯定会先递归到底，然后呢，我们就可以开始统计答案了，如果是轻儿子，我们在统计完这个结点向下的答案之后，我们需要更新该结点，清空它的贡献。于是，根据调和级数的定义，我们能确定最后的复杂度是O(N*log(N))级别的。

#include <iostream>
#include <cstdio>
#include <cmath>
#include <string>
#include <cstring>
#include <algorithm>
#include <limits>
#include <vector>
#include <stack>
#include <queue>
#include <set>
#include <map>
#include <unordered_map>
#include <unordered_set>
#define _ABS(x, y) ( x > y ? (x - y) : (y - x) )
#define lowbit(x) ( x&( -x) )
#define pi 3.141592653589793
#define e 2.718281828459045
#define INF 0x3f3f3f3f
#define efs 1e-7
#define HalF (l + r)>>1
#define lsn rt<<1
#define rsn rt<<1|1
#define Lson lsn, l, mid
#define Rson rsn, mid+1, r
#define QL Lson, ql, qr
#define QR Rson, ql, qr
#define myself rt, l, r
#define MP(a, b) make_pair(a, b)
using namespace std;
typedef unsigned long long ull;
typedef unsigned int uit;
typedef long long ll;
const int maxN = 1e5 + 7;
int N, head[maxN], cnt, col[maxN], num[maxN];
struct Eddge
{int nex, to;Eddge(int a=-1, int b=0):nex(a), to(b) {}
}edge[maxN << 1];
inline void addEddge(int u, int v)
{edge[cnt] = Eddge(head[u], v);head[u] = cnt++;
}
inline void _add(int u, int v) { addEddge(u, v); addEddge(v, u); }
int Wson[maxN], siz[maxN];
void W_dfs(int u, int fa)
{siz[u] = 1;int maxx = 0;for(int i=head[u], v; ~i; i=edge[i].nex){v = edge[i].to;if(v == fa) continue;W_dfs(v, u);siz[u] += siz[v];if(siz[v] > maxx){maxx = siz[v];Wson[u] = v;}}
}
ll sum = 0, ans[maxN] = {0};
int Mx, son;    //出现次数，重儿子的贡献不统计
inline void add(int u, int fa, int val)
{num[col[u]] += val; //删或者加if(num[col[u]] > Mx) { Mx = num[col[u]]; sum = col[u]; }else if(num[col[u]] == Mx) sum += col[u];for(int i=head[u], v; ~i; i=edge[i].nex){v = edge[i].to;if(v == fa || v == son) continue;add(v, u, val);}
}
void dfs(int u, int fa, int op)
{for(int i=head[u], v; ~i; i=edge[i].nex){v = edge[i].to;if(v == fa || v == Wson[u]) continue;dfs(v, u, 0);   //走的是轻儿子}if(Wson[u]) { dfs(Wson[u], u, 1); son = Wson[u]; } //统计重儿子的贡献，不消除影响add(u, fa, 1); son = 0; //暴力统计所有轻儿子的贡献ans[u] = sum;if(!op) { add(u, fa, -1); sum = 0; Mx = 0; }    //因为重儿子的答案是最后才统计的，直接在重链往上，不消除
}
inline void init()
{cnt = 0;for(int i=1; i<=N; i++) head[i] = -1;
}
int main()
{scanf("%d", &N);init();for(int i=1; i<=N; i++) scanf("%d", &col[i]);for(int i=1, u, v; i < N; i++){scanf("%d%d", &u, &v);_add(u, v);}W_dfs(1, 0);dfs(1, 0, 0);for(int i=1; i<=N; i++) printf("%lld%c", ans[i], i == N ? 'n' : ' ');return 0;
}