Rainbow Roads(gym101617G)(DFS序,差分)

题意：

题目给出一棵树，树上每条边都有对应颜色，我们定义一条简单路径上没有连续出现两个相同的颜色，则称这条简单路径为彩虹路。如果一个节点通往其他所有节点的路径都是彩虹路，则称这个节点为超级节点。问这棵树上有多少个超级节点并输出超级节点。

---

想法：

• 首先考虑哪些点是不能成立的。
• 1)对于一个节点来说，如果它到它的两个子节点的边的颜色相同，那么这两个子节点的子树全部不成立(包括两个子节点)。
• 2)对于一个节点来说，如果它到它的一个子节点的边的颜色和它到它的父亲节点的边的颜色相同，那么除了这个节点和这个节点另外的子节点外，其他都不成立。

---

• 下面考虑怎么统计。
• 我们可以通过 \(DFS\) 求出这棵树的 \(dfs\)序 ( \(dfn[i]\) )、所有节点的父节点 ( \(fa[i]\)）、所有节点的子树大小( \(son[i]\) )。
• 然后通过 \(dfs\)序，那么我们就可以通过差分数组，标记每个不成立区间，然后最后就求出成立区间。
• 接下来就是枚举颜色相同且相邻的两条边，去考虑上述两种情况（以下区间为左闭右开区间）。
• 对于第一种情况，设两个子节点分别是 \(x,y\)，那么我们需要标记的区间是 $dfn[x]\sim dfn[x]+son[x] $和 $dfn[y]\sim dfn[y]+son[y] $，也就是左边界+1，右边界-1。
• 对于第二种情况，我们设 \(x\) 为当前节点的父节点， \(y\) 为当前节点的子节点，需要标记的区间就是 \(1\sim dfn[fa[i]]\)、\(dfn[i]+son[i]\sim n\)、\(dfn[x]\sim dfn[x]+son[x]\)。
• 最后通过这个差分数组求出 \(vis\) 数组， \(vis[i]=0\) 即使满足的节点。

---

代码：

#include<bits/stdc++.h>
#define ll long long
#define mod 1000000007
#define pdd pair<double,double>
#define mem(a,x) memset(a,x,sizeof(a))
#define IO ios::sync_with_stdio(false);cin.tie(0)
using namespace std;
const long long INF = 1e18+5;
const int inf = 0x3f3f3f3f;  
const double eps=1e-6;       
const int maxn=80005;    
const double pi=acos(-1);   
inline char nc() {static char buf[1000000],*p1=buf,*p2=buf;return p1==p2&&(p2=(p1=buf)+fread(buf,1,1000000,stdin),p1==p2)?EOF:*p1++;}
inline void read(int &sum) {char ch=nc();sum=0;while(!(ch>='0'&&ch<='9')) ch=nc();while(ch>='0'&&ch<='9') sum=(sum<<3)+(sum<<1)+(ch-48),ch=nc();}

struct edge{
    int v;
    int color;
};
vector<edge>g[maxn];
bool cmp(edge aa,edge bb)
{
    return aa.color<bb.color;
}
int dfn[maxn],fa[maxn],son[maxn],dir[maxn];
int arr[maxn],vis[maxn];
vector<int>ans;
int n;
int cnt=0;


void dfs(int now,int pre)
{
    dfn[now]=++cnt;
    son[now]=1;
    fa[now]=pre;
    dir[cnt]=now;
    for(auto it:g[now]){
        int v=it.v;
        if(v==pre)continue;
        dfs(v,now);
        son[now]+=son[v];
    }
}

int main()
{
    scanf("%d",&n);
    for(int i=1;i<n;i++){
        int u,v,color;
        scanf("%d%d%d",&u,&v,&color);
        g[u].push_back((edge){v,color});
        g[v].push_back((edge){u,color});
    }
    dfs(1,0);
    for(int i=1;i<=n;i++){
        sort(g[i].begin(),g[i].end(),cmp);
        int siz=g[i].size();
        for(int j=0;j<siz-1;j++){
            if(g[i][j].color==g[i][j+1].color){
                int x=g[i][j].v,y=g[i][j+1].v;
                if(fa[x]==i&&fa[y]==i){
                    arr[dfn[x]]++;
                    arr[dfn[x]+son[x]]--;
                    arr[dfn[y]]++;
                    arr[dfn[y]+son[y]]--;
                }else if(fa[x]==i&&fa[i]==y){
                    arr[1]++;
                    arr[dfn[i]]--;
                    arr[dfn[i]+son[i]]++;
                    arr[dfn[x]]++;
                    arr[dfn[x]+son[x]]--;
                }else if(fa[y]==i&&fa[i]==x){
                    arr[1]++;
                    arr[dfn[i]]--;
                    arr[dfn[i]+son[i]]++;
                    arr[dfn[y]]++;
                    arr[dfn[y]+son[y]]--;
                }
            }
        }
    }
    for(int i=1;i<=n;i++){
        vis[i]+=vis[i-1]+arr[i];
        if(vis[i]==0)ans.push_back(dir[i]);
    }
    sort(ans.begin(),ans.end());
    cout<<ans.size()<<endl;
    for(auto it:ans){
        cout<<it<<endl;
    }
} 

/*
8 
1 3 1
2 3 1
3 4 3
4 5 4
5 6 3
6 7 2
6 8 2
*/