dfs序 Rainbow Roads

Rainbow Roads

题目连接：https://codeforces.com/group/mey3UXMrvB/contest/254921

 

这种题，，，晚上睡醒写完，都不带调试的过了，下午太困了，，连思路都没有

#include<bits/stdc++.h>
using namespace std;
const int M=1e5+4;
struct Node{int to,nex,col;}e[M<<1];
struct node{int x,y;};
int head[M<<1],cnt,n,dfn[M],dep[M],res[M],siz[M],rk[M];
stack<node> s;
void dfs1(int x,int f){
    dfn[x]=++cnt;
    dep[x]=dep[f]+1;
    siz[x]=1;
    rk[cnt]=x;
    for(int to,i=head[x];~i;i=e[i].nex)
        if(e[i].to!=f){
            dfs1(e[i].to,x);
            siz[x]+=siz[e[i].to];
        }
}
void dfs2(int x,int f){
    while(!s.empty())s.pop();
    node t;
    for(int i=head[x];~i;i=e[i].nex)
        for(int j=e[i].nex;~j;j=e[j].nex){
            if(e[j].col==e[i].col){
                t.x=e[i].to;t.y=e[j].to;
                if(dfn[t.x]>dfn[t.y])swap(t.x,t.y);
                s.push(t);
            }
        }
    while(!s.empty()){
        t=s.top();s.pop();
        if(dep[t.x]!=dep[t.y]){
            res[1]++;res[dfn[x]]--;res[dfn[x]+siz[x]]++;
            res[dfn[t.y]]++;res[dfn[t.y]+siz[t.y]]--;
        }else{
            res[dfn[t.x]]++;res[dfn[t.x]+siz[t.x]]--;
            res[dfn[t.y]]++;res[dfn[t.y]+siz[t.y]]--;
        }
    }
    for(int i=head[x];~i;i=e[i].nex)
        if(e[i].to!=f)dfs2(e[i].to,x);
    return ;
}
int main(){
    memset(head,-1,sizeof(head));
    scanf("%d",&n);
    for(int x,y,col,i=1;i<n;i++){
        scanf("%d%d%d",&x,&y,&col);
        e[cnt].to=y;e[cnt].col=col;e[cnt].nex=head[x];head[x]=cnt++;
        e[cnt].to=x;e[cnt].col=col;e[cnt].nex=head[y];head[y]=cnt++;
    }cnt=0;
    dfs1(1,-1);
    dfs2(1,-1);
    res[0]=0;
    vector<int> ans;
    for(int i=1;i<=n;i++){
        res[i]+=res[i-1];
        if(res[i]==0){
            res[0]++;
            ans.push_back(rk[i]);
        }
    }
    sort(ans.begin(),ans.end());
    printf("%d\n",res[0]);
    for(auto x:ans)printf("%d\n",x);
    return 0;
}