\(\text{Description}\)
\(\text{Solution}\)
\(\text{dsu on tree}\)
好妙啊。
显然统计了 \(u\) 为根子树后要清空才能统计它的兄弟，我们尝试加一些小小的优化。
预处理出轻重儿子，由于如果一个儿子是最后遍历的就不用清空，直观地我们选择重儿子最后遍历。
如何证明时间复杂度？考虑一个点到根的链上最多只有 \(\log n\) 条轻边。因为一条轻边意味着节点的父亲的 \(\rm siz\) 大于此节点 \(\rm siz\) 的两倍，所以一条轻边 \(\rm siz\) 就会乘以 \(2\)，最多只能乘 \(\log n\) 次。
那么，由于重儿子 \(\rm son\) 不被清空，所以 \(\rm son\) 的 父亲 就不需要再次遍历 \(\rm son\) 以获得答案，而它的父亲的轻儿子是需要 全部遍历 的。那么可以发现对于点 \(x\) 而言，当它的某个祖先 \(y\) 向上连接了一条重边，那么 \(y\) 的父亲 \(z\) 就不会再遍历 \(y\) 这一棵子树，也就是说 \(x\) 只会被遍历 "到根的链上的轻边数" 次，时间复杂度 \(\mathcal O(n\log n)\)。
线段树合并
虽然理论上这也是 \(\mathcal O(n\log n)\) 的，但写出来不仅内存大而且多了一秒有余…
对于每个点开一个权值线段树来维护子树颜色及其个数。先遍历子树，将自己和儿子 \(\rm merge\) 一下就行了。其他的复杂度都是单次 \(\mathcal O(\log n)\) 的。
合并的复杂度可以这样计算：一开始有 \(n\) 棵树，每棵树只有一个位置有权值，节点总数是 \(n\log n\) 的。合并复杂度是 merge() 递归的次数，我们可以将其分成两种情况 —— 还能再调用函数与不能再调用函数。因为满二叉树有叶节点个数与非叶节点个数同阶的性质，这两种情况的次数也是同阶的，我们只用计算 "还能再调用函数" 的次数。可以发现，每发生一次 "还能再调用函数" 就会有一个 \(y\) 点被扔掉，所以它的次数可以看作总点数，所以总复杂度是 \(\mathcal O(n\log n)\) 的。
\(\text{Code}\)
\(\text{dsu on tree}\)
#include <cstdio>
#include <vector>
using namespace std;
typedef long long ll;
const int maxn=1e5+5;
int n,c[maxn],f[maxn],tot[maxn],siz[maxn],son[maxn],maxx;
ll ans[maxn],sum;
bool vis[maxn];
vector <int> e[maxn];
void dfs1(int u,int fa) {
f[u]=fa,siz[u]=1;
for(int i=0;i<e[u].size();++i) {
if(e[u][i]==fa) continue;
dfs1(e[u][i],u);
siz[u]+=siz[e[u][i]];
if(siz[e[u][i]]>siz[son[u]]) son[u]=e[u][i];
}
}
void calc(int u,int k) {
tot[c[u]]+=k;
if(tot[c[u]]>maxx) maxx=tot[c[u]],sum=c[u];
else if(tot[c[u]]==maxx) sum+=c[u];
for(int i=0;i<e[u].size();++i)
if(!vis[e[u][i]]&&e[u][i]!=f[u])
calc(e[u][i],k);
}
void dfs2(int u,int Keep) {
for(int i=0;i<e[u].size();++i) if(e[u][i]!=son[u]&&e[u][i]!=f[u]) dfs2(e[u][i],0);
if(son[u]) vis[son[u]]=1,dfs2(son[u],1);
calc(u,1); ans[u]=sum;
if(son[u]) vis[son[u]]=0;
if(!Keep) calc(u,-1),sum=maxx=0;
}
int main() {
scanf("%d",&n);
for(int i=1;i<=n;++i) scanf("%d",&c[i]);
for(int i=1;i<n;++i) {
int x,y;
scanf("%d %d",&x,&y);
e[x].push_back(y),e[y].push_back(x);
}
dfs1(1,0); dfs2(1,1);
for(int i=1;i<=n;++i) printf("%lld ",ans[i]); puts("");
return 0;
}
线段树合并
#include <cstdio>
#define print(x,y) write(x),putchar(y)
template <class T>
inline T read(const T sample) {
T x=0; char s; bool f=0;
while((s=getchar())>'9' or s<'0')
f|=(s=='-');
while(s>='0' and s<='9')
x=(x<<1)+(x<<3)+(s^48),
s=getchar();
return f?-x:x;
}
template <class T>
inline void write(const T x) {
if(x<0) {
putchar('-'),write(-x);
return;
}
if(x>9) write(x/10);
putchar(x%10^48);
}
#include <vector>
using namespace std;
typedef long long ll;
const int maxn=1e5+5;
int n,a[maxn],idx;
int rt[maxn];
vector <int> e[maxn];
ll Ans[maxn];
struct node {
int mx,ls,rs; ll ans;
} t[maxn*80];
void pushUp(int o) {
int l=t[o].ls,r=t[o].rs;
if(t[l].mx>t[r].mx) {
t[o].mx=t[l].mx;
t[o].ans=t[l].ans;
}
else if(t[l].mx<t[r].mx) {
t[o].mx=t[r].mx;
t[o].ans=t[r].ans;
}
else {
t[o].mx=t[l].mx;
t[o].ans=t[l].ans+t[r].ans;
}
}
int merge(int x,int y,int l,int r) {
if(!x or !y) return x|y;
if(l==r) {
t[x].ans=l;
t[x].mx+=t[y].mx;
return x;
}
int mid=l+r>>1;
t[x].ls=merge(t[x].ls,t[y].ls,l,mid);
t[x].rs=merge(t[x].rs,t[y].rs,mid+1,r);
pushUp(x);
return x;
}
void ins(int &o,int l,int r,int p,int k) {
if(!o) o=++idx;
if(l==r) {
t[o].ans=l;
t[o].mx+=k;
return;
}
int mid=l+r>>1;
if(p<=mid)
ins(t[o].ls,l,mid,p,k);
else ins(t[o].rs,mid+1,r,p,k);
pushUp(o);
}
void dfs(int u,int fa) {
for(auto v:e[u]) {
if(v==fa) continue;
dfs(v,u);
rt[u]=merge(rt[u],rt[v],1,1e5);
}
ins(rt[u],1,1e5,a[u],1);
Ans[u]=t[rt[u]].ans;
}
int main() {
n=read(9);
for(int i=1;i<=n;++i)
a[i]=read(9);
for(int i=1;i<n;++i) {
int u=read(9),v=read(9);
e[u].push_back(v);
e[v].push_back(u);
}
dfs(1,0);
for(int i=1;i<=n;++i)
print(Ans[i],' ');
puts("");
return 0;
}
