J 树上行走

  • 分析

    似乎有一种方法是并查集。但是这里提供一种新的思路。先画个图。
    图片说明

    首先我们可以在第一次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;
}