图论训练之九

https://www.luogu.org/problem/P5588?contestId=22026

分析：

首先同一颜色在树上连边一定是一条链，

如果没有这个颜色的结点答案则为n*(n-1)/2

如果该颜色只有一个节点，那么就是统计一个这节点为根的所有子树大小互乘累加

下面重点讨论

一个颜色有两个以上节点的情况，那么可能会出现两种情况：

一：所有这样的节点都在一条链上，也就是说这个颜色有两个端点。

二：所有这样的颜色不在一条链上，也就是说有三个及以上端点

那么显而易见的，对于第2种情况一定不存在一种合法的解，输出 0即可。

所以重点是解决第1种情况

我们考虑对于一个颜色，什么时候可能会是端点。

我们记一个数组 cnt表示颜色 c到目前为止出现的次数。

那么对于每一次cnt改变，我们定义一个变量 flag++，只要最终结果 flag=1 就极可能是个端点。

还有另一种可能，就是进入当前节点时 cnt[c]已经有值，或者当前节点不为当前颜色的最后一个节点，那么也要使 flag++。

具体再解释下 flag 可能会更好理解：flag 代表当前节点下子树颜色也为 c 的个数。

记得最后处理完上面的步骤还要让 cnt[c]++

只要最后 flag=1那么即为端点。

第一次进入直接让当前颜色的节点指针赋给 nos (记录当前颜色从哪开始)

然后答案即为最两端的节点的子树大小之积。

code ：

#include<bits/stdc++.h>
using namespace std;
const int N = 1e6 + 10;
int head[N], Next[N << 1], ver[N << 1];
void add(int x, int y) {
    static int cnt = 0;
    ver[++cnt] = y, Next[cnt] = head[x], head[x] = cnt;
    ver[++cnt] = x, Next[cnt] = head[y], head[y] = cnt;
}
int color[N], tot[N], cnt[N], n;
int enos[N], size[N], nos[N];
long long ans1[N], ans2[N];
inline void dfs(int x, int fa) {
    int c = color[x], k = cnt[c];
    int flag = 0, pos = 0;
    size[x] = 1;
    for(int i = head[x]; i; i = Next[i]) {
        int y = ver[i];
        if(y == fa) continue;
        int lastans = cnt[c];
        dfs(y, x);
        ans1[x] += 1LL * size[x] * size[y];
        size[x] += size[y];
        if(lastans != cnt[c]) flag++, pos = y;
    }
    ans1[x] += 1LL * size[x] * (n - size[x]);
    if(k || cnt[c] != tot[c] - 1) flag++;
    cnt[c]++;
    if(flag == 1) {
        if(!enos[c]) nos[c] = x;
        else {
            int p = pos ? n - size[pos] : size[x];
            ans2[c] = 1LL * size[nos[c]] * p;
        } enos[c]++;    
    }
}

int main() {
    scanf("%d", &n);
    for(int i = 1; i <= n; i++) {
        scanf("%d", color + i);
        tot[color[i]]++;
        nos[color[i]] = i;
    }
    for(int i = 1, x, y; i < n; i++) {
        scanf("%d %d", &x, &y);
        add(x, y);
    } 
    dfs(1, 0);
    for(int i = 1; i <= n; i++) {
        if(tot[i] == 0) 
            printf("%lld\n", 1LL * n * (n - 1) >> 1);
        else if(tot[i] == 1)
            printf("%lld\n", ans1[nos[i]]);
        else if(enos[i] == 2)
            printf("%lld\n", ans2[i]);
        else puts("0");
    }
    return 0;
}