E. Count Paths

原题链接

题解

前言：这题可以只调用一遍dfs。

首先，以颜色为color_u的u为根结点的子树内，颜色与u颜色相同的结点 不能与u的其余子树中颜色为color_u的结点相连接。

我们需要一个num数组，num[i]表示当前结点 i ，有多少个结点可以与他连接。

接下来，我们任取一个结点为根结点去跑dfs。假设当前在 i 结点，在每次遍历一条子路径时，将num[i]取1，在退出当前结点时num[i]恢复并+1。

具体思想和细节可以看code

code

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N=2e5+5;
const int M=4e5+5;
int head[N],Next[M],to[M],cnt=1,color[N],vis[N];
ll ans=0,num[N];
void Init(int n){
    for (int i=1;i<=n;i++){
        head[i]=0;
        num[i]=0;
        vis[i]=0;
    }
    cnt=1;
    ans=0ll;
}
void build(int from,int to_){
    Next[cnt]=head[from];
    to[cnt]=to_;
    head[from]=cnt++;
}
void dfs(int tree){
    if (vis[tree]) return;
    vis[tree]=1;
    int c=color[tree];
    ll x=num[c];
    ans+=x;
    for (int i=head[tree];i>0;i=Next[i]){
        num[c]=1;   //每次遍历子树都要将根结点的颜色变 1 
        dfs(to[i]);
    }
    num[c]=x+1ll;
    vis[tree]=0;
}
int main(){
    ios::sync_with_stdio(false);
    int t;
    cin>>t;
    while (t--){
        int n;
        cin>>n;
        Init(n);
        for (int i=1;i<=n;i++) cin>>color[i];
        for (int i=1,u,v;i<n;i++){
            cin>>u>>v;
            build(u,v);
            build(v,u);
        }
        dfs(1);
        cout<<ans<<endl;
    }
    return 0;
}