分析
似乎有一种方法是并查集。但是这里提供一种新的思路。先画个图。
首先我们可以在第一次dfs的时候求出对于以1为根的某棵子树内部有多少相连的颜色一样的点
inline void dfs(int u,int v) { sum[u]=1; for (int i=h[u];~i;i=nex[i]) { int j=ver[i]; if(j==v) continue; dfs(j,u); if(a[u]==a[j]) sum[u]+=sum[j]; } }
接下来我们考虑怎么求到以某一个点为根,在不同的类型的点消失后,他能直接相连的点数(设为k)。
也就是说,我们要想个办法使得在遍历树的时候能够从上往下更新。假设到了一个点u,v是它的父节点,那么k[v]=sum[u]+(其他子树的贡献),也就是说,如果对于一个点往上走能到达的点数其实就是k[v]-sum[u]
inline void dow(int u,int v) { k[u]=sum[u]; if(a[u]==a[v]) { int rem=k[v]-sum[u]; k[u]+=rem; } for (int i=h[u];~i;i=nex[i]) { int j=ver[i]; if(j==v) continue; dow(j,u); } }
剩下的就是统计结果了
代码
#include<bits/stdc++.h> using namespace std; const int N=2e5+10; int n,tot,cnt; int a[N],k[N],sum[N]; int ans[N],h[N],nex[N<<1],ver[N<<1]; inline void add(int x,int y) { nex[tot]=h[x]; ver[tot]=y; h[x]=tot++; } inline void dfs(int u,int v) { sum[u]=1; for (int i=h[u];~i;i=nex[i]) { int j=ver[i]; if(j==v) continue; dfs(j,u); if(a[u]==a[j]) sum[u]+=sum[j]; } } inline void dow(int u,int v) { k[u]=sum[u]; if(a[u]==a[v]) { int rem=k[v]-sum[u]; k[u]+=rem; } for (int i=h[u];~i;i=nex[i]) { int j=ver[i]; if(j==v) continue; dow(j,u); } } int main() { memset(h,-1,sizeof(h)); scanf("%d",&n); for (int i=1;i<=n;i++) scanf("%d",&a[i]); for (int i=1;i<n;i++) { int u,v; scanf("%d%d",&u,&v); add(u,v),add(v,u); } dfs(1,0); a[0]=(a[1]^1); sum[0]=sum[1];k[0]=sum[1]; dow(1,0); int ma=k[1];ans[cnt=1]=1; for (int i=2;i<=n;i++) { if(ma<k[i]) ma=k[i],ans[cnt=1]=i; else if(ma==k[i]) ans[++cnt]=i; } printf("%d\n",cnt); for (int i=1;i<=cnt;i++) printf("%d ",ans[i]); return 0; }