【JZOJ 5438】【NOIP2017提高A组集训10.31】Tree

Description

Solution

有一种方便的做法，
我们要保证儿子与父亲一样嘛，最后再考虑根，
对于一条边：x->y，如果x,y的颜色不一样，就在y上标记，

复杂度：$O(n)$

Cdeo

#include <cstdio>
#include <algorithm>
#define fo(i,a,b) for(int i=a;i<=b;i++)
#define fod(i,a,b) for(int i=a;i>=b;i--)
#define efo(i,q) for(int i=A[q];i;i=B[i][0])
#define min(q,w) ((q)>(w)?(w):(q))
#define max(q,w) ((q)<(w)?(w):(q))
using namespace std;
const int N=500500;
int read(int &n)
{
    char ch=' ';int q=0,w=1;
    for(;(ch!='-')&&((ch<'0')||(ch>'9'));ch=getchar());
    if(ch=='-')w=-1,ch=getchar();
    for(;ch>='0' && ch<='9';ch=getchar())q=q*10+ch-48;n=q*w;return n;
}
int m,n,ans;
int B[2*N][2],A[N],B0;
int Ans[N],b[N];
int d[N][3];
void link(int q,int w)
{
    B[++B0][0]=A[q],A[q]=B0,B[B0][1]=w;
    B[++B0][0]=A[w],A[w]=B0,B[B0][1]=q;
}
void bfs()
{
    int S=1,T=1;
    d[1][0]=1;
    d[1][1]=0;
    d[1][2]=0;
    for(;S<=T;S++)
    {
        int q=d[S][0];
        int Cl=d[S][2];
        b[q]=b[q]^Cl;
        if(b[q])Ans[++ans]=q,Cl=!Cl,b[q]=0;
        efo(i,q)if(B[i][1]!=d[S][1])
        {
            d[++T][0]=B[i][1];
            d[T][1]=q;
            d[T][2]=Cl;
        }
    }
}
int main()
{
    freopen("tree.in","r",stdin);
    freopen("tree.out","w",stdout);
    int q,w;
    read(n);
    fo(i,1,n)read(b[i]);
    fo(i,1,n-1)read(q),read(w),link(q,w);
    bfs();
    sort(Ans+1,Ans+1+ans);
    fo(i,1,ans)printf("%d ",Ans[i]);
    printf("\n");
    return 0;
}