Colorful Tree（逆向思维）

Problem Description
There is a tree with n nodes, each of which has a type of color represented by an integer, where the color of node i is ci.

The path between each two different nodes is unique, of which we define the value as the number of different colors appearing in it.

Calculate the sum of values of all paths on the tree that has n(n−1)2 paths in total.

Input
The input contains multiple test cases.

For each test case, the first line contains one positive integers n, indicating the number of node. (2≤n≤200000)

Next line contains n integers where the i-th integer represents ci, the color of node i. (1≤ci≤n)

Each of the next n−1 lines contains two positive integers x,y (1≤x,y≤n,x≠y), meaning an edge between node x and node y.

It is guaranteed that these edges form a tree.

Output
For each test case, output “Case #x: y” in one line (without quotes), where x indicates the case number starting from 1 and y denotes the answer of corresponding case.

Sample Input
3
1 2 1
1 2
2 3
6
1 2 1 3 2 1
1 2
1 3
2 4
2 5
3 6

Sample Output
Case #1: 6
Case #2: 29
题目大意：给你一个包含n个节点的树，用一个数字代表一种颜色，树上的路径权值为在这条路径上包含的颜 色数量，这棵树总共有n*(n-1)/2条路径。求这棵树上所有路径的总权值。
解题思路：首先求所有路径上包含颜色的总数量等价于求每种颜色有多少条路径经过，又等价于 颜色总数 * 路径总数 - 每种颜色没有经过的路径的总和。显然，这里颜色总数饿路径总数都很好求，我们现在要做的就是求每种颜色有多少条路径没有经过。这里有一个小技巧，就是每种颜色不经过的路数=每种颜色不经的点数 * 不经过的点数-1/2（即：count*(count-1)/2）；每种颜色不经过点的个数可以分为两种
1.与这个颜色的节点有关联但颜色不同的点
2.与这个颜色的节点无关的且颜色不同的点
如第二组测试样例，所构成的树为

这里我把1用红色表示，2用黄色表示，3用绿色表示
颜色总数 * 路径总数 = 3 * （6 * (6-1)）/2 = 45;
未经过红色的路径总数为：3 * 2/ 2 = 3；
未经过黄色的路径总数为：1 * 0 / 2 + 3 * 2 / 2 = 3;
未经过黄色的路径总数为：5 * 4 / 2 = 10;
所以树的权值为 45- 3 -3-10 = 29；
思路就到这里了，下边是代码实现，这里提供两个版本，一个是自己写的领接链表，另一个是用vector表示的领接链表
第一种：

#include<bits/stdc++.h>
#define maxn 200005
#define inf 0x3f3f3f3f
using namespace std;
typedef long long LL;
const double eps = 1e-8;
const int mod = 1e9 + 7;
int c[maxn], head[maxn], len, sum[maxn], size[maxn], vis[maxn];
//c数组用来记录输入节点的颜色，vis用来标记有几种颜色，size[i]用来记录以i节点为根节点的节点数目
//head用来记录邻接链表的头结点，sum[i]用来表示以i为根节点的子树中颜色与根节点颜色相同的子树
LL d;//第一种情况下未经过路径数
struct node{
    int v, next;
}p[maxn << 1];//双向的最大节点数*2
void addedge(int u, int v){//向邻接链表里添加边
    p[len].v = v;
    p[len].next = head[u];
    head[u] = len;
    len++;
}
void dfs(int x, int fa){//x表示当前节点，fa表示为它的前驱节点，1的前驱节点为0对树进行深度优先搜索，搜索顺序为1 3 6 2 5 2 找到每个颜色有哪些路没有 经过，并用b累加起来
    LL pre = sum[c[x]];///记录当前节点以前与当前节点颜色相同的节点数
    //sum[c[x]]++;
    size[x] = 1;///当前节点x为根节点的节点为1
    int add = 0;///记录与从x颜色节点有关的节点个数
    for (int i = head[x]; i!=-1; i = p[i].next){
        if (p[i].v == fa)
            continue;
        dfs(p[i].v, x);
        size[x] += size[p[i].v];//更新以X为根节点的节点个数
        LL count = size[p[i].v] - sum[c[x]] + pre;///count表示与当前节点颜色不同的子节点个数
        pre = sum[c[x]];///更新pre，加上子树中x出现的次数
        add += count;
        d += count*(count - 1) >> 1;
    }
    sum[c[x]] += add + 1;
}
int main(){
    int n, tcase = 1;
    while (~scanf("%d", &n)){
        memset(head, -1, sizeof(head));
        memset(sum, 0, sizeof(sum));
        memset(vis, 0, sizeof(vis));
        d = len = 0;
        LL number = 0;
        for (int i = 1; i <= n; i++){
            scanf("%d", &c[i]);
            if (!vis[c[i]]){
                vis[c[i]] = 1;
                number++;
            }
        }
        for (int i = 1; i < n; i++){
            int u, v;
            scanf("%d%d", &u, &v);
            addedge(u, v);
            addedge(v, u);
        }
        dfs(1, 0);
        LL ans = (number*(n - 1)*n >> 1) - d;//利用反面来算，所有的颜色数*路的总数减去第一种情况下未经过路径数
        for (int i = 1; i <= n; i++){
            if (vis[i] && i != c[1]){
                LL count = n - sum[i];
                ans -= count*(count - 1) >> 1;///减去第二种情况下未经过路径数
            }
        }
        printf("Case #%d: %lld\n", tcase++, ans);
    }
    return 0;
}

第二种：

#include <bits/stdc++.h>
using namespace std;
const int N = 4e5+100;
typedef long long ll;
int c[N],vis[N];///c数组表示第i个节点的颜色ci;vis为标记数组
vector<int> e[N];///e数组用来存储树的信息
ll sum[N],size[N];///sum数组模拟根节点i的子树相同颜色的个数；size[i]表示的是以i为根的点的个数，sun[i]-size[i]=与根节点不同的颜色的个数（sum，size表示是在统一课子树的）
ll ans;///ans表示不经过颜色为i的路径总条数
void dfs(int x,int y)///x表示当前节点，y表示其前一个节点
{

    size[x]=1;  //自身
    sum[c[x]]++;   //加上自身
    ll pre=sum[c[x]];   //pre表示跟节点之前相同颜色的数量也就是最高的数到这个根节点之间的
    for(int i=0; i<e[x].size(); i++)
    {

        if(e[x][i]==y)
            continue;
        dfs(e[x][i],x);
        size[x]+=size[e[x][i]];/// 以x为根的树的点的总个数，当前的这个点还要加上他的子数上的点
        ll count=size[e[x][i]]-(sum[c[x]]-pre);///count表示与当前节点颜色不同的子节点个数
        ans=ans+(1LL*count*(count-1))/2;
        sum[c[x]]+=count;///将算入路径中的颜色模拟成节点x的颜色。以防更高根节点有相同颜色重复计算。这里注意，如果子树中无x颜色的点的个数为1时，也是构不成路径，但是也要算在sum[x]里面，因为这样的节点在整棵树中也是无法构成路径的
        pre=sum[c[x]];   //这里pre比之前多了子树部分的，防止其他新子树的影响
    }
}

int main()
{
    int n,cas=1;
    while(scanf("%d",&n)!=EOF)
    {
        int num=0;
        ans=0;
       memset(sum,0,sizeof(sum));
        memset(vis,0,sizeof(vis));
        for(int i=1; i<=n; i++)
        {
            e[i].clear();
            scanf("%d",&c[i]);
            if(vis[c[i]]==0)
            {
                vis[c[i]]=1;///标记ci颜色
                num++;
            }
        }
        for(int i=1; i<n; i++)
        {
            int u,v;
            scanf("%d%d",&u,&v);
            e[u].push_back(v);
            e[v].push_back(u);
        }
        dfs(1,0);   //求与根节点相关点要减去路径的路数
        ll ANS = 1LL*num*((1LL)*n*(n-1))/2;/// 这里要算的就是整棵树中所有颜色都经过每一条路径的所有和
        for(int i=1; i<=n; i++)    ///求与根节点无关的点要减去的路径数
        {
            if(vis[i])
            {
                ll ct=n-sum[i];///ct表示模拟过后树中没有颜色i的节点个数
                ans+=ct*(ct-1)/2;
            }
        }
        printf("Case #%d: %lld\n", cas++, ANS-ans);
    }
}