Atcoder - 163F path pass i

传送门：Atcoder - 163F

思路分析

对于每一个点，经过这个点的路径数量很容易计算，但是如果路径上有颜色重复的点的话
答案就会重复，考虑如果去掉重复的答案
计算一个点的时候，之前的出现过的相同颜色的点的子树就是重复的答案，将这个数量删去
每次计算完一颗子树后，将这个子树根颜色的被影响值加上这棵子树的大小
在计算一棵子树内的答案的时候，子树根的颜色被影响值=n-子树大小，这是显然的

样例输入

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

样例输出

18
15
0
14
23
0
23
0

AC代码

#include <functional>
#include <algorithm>
#include <iostream>
#include <cstring>
#include <cstdlib>
#include <iomanip>
#include <vector>
#include <string>
#include <cstdio>
#include <chrono>
#include <random>
#include <queue>
#include <stack>
#include <cmath>
#include <map>
#include <set>
#if __cplusplus >= 201103L
#include <unordered_map>
#include <unordered_set>
#endif
#define ls x<<1
#define rs x<<1|1
#define fi first
#define se second
#define ll long long
#define pb push_back
#define mp make_pair
#define fun function
#define sz(x) x.size()
#define lowbit(x) x&(-x)
#define all(x) x.begin(),x.end()
#define mem(a,b) memset(a,b,sizeof(a))
#define IOS ios_base::sync_with_stdio(0); cin.tie(0);cout.tie(0);
#define int long long
using namespace std;

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

const int INF = 0x3f3f3f3f;
const int N=1e6+10;

int a[N],sum[N],siz[N],ans[N],n;
vector<int>e[N];

void dfs1(int u,int fa) {
    siz[u]=1;
    for(auto v:e[u]) {
        if(v==fa) continue;
        dfs1(v,u);
        siz[u]+=siz[v];
    }
}

void dfs2(int u,int fa) {
    ans[a[u]]+=siz[u]*(n-siz[u]-sum[a[u]]+1);
    sum[a[u]]+=siz[u];
    int ss=sum[a[u]];
    int pre=0;
    for(auto v:e[u]) {
        if(v==fa) continue;
        sum[a[u]]=n-siz[v];
        dfs2(v,u);
        ans[a[u]]+=pre*siz[v];
        pre+=siz[v];
    }
    sum[a[u]]=ss;
}

signed main() {
    IOS;
#ifdef xiaofan
    freopen("1.in","r",stdin);
    freopen("1.out","w",stdout);
#endif

    cin>>n;
    for(int i=1; i<=n; i++) cin>>a[i];
    for(int i=1; i<n; i++) {
        int u,v;
        cin>>u>>v;
        e[u].push_back(v);
        e[v].push_back(u);
    }
    dfs1(1,1);
    dfs2(1,1);
    for(int i=1; i<=n; i++) cout<<ans[i]<<endl;



    return 0;
}