HDU 6035-Colorful Tree(dfs+树)

题目链接 http://acm.hdu.edu.cn/showproblem.php?pid=6035
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）；
这样我们就要求每种颜色不经过的点的个数，而每种颜色不经过点的个数

可以分为两种：

第一种为与根节点有关联但不颜色不同的点：
如下图以节点2为根节点，节点4就是与2相关联但 不相同的点。节点1为根节点，节点2，4， 5。

第二中为与根节点无关的且颜色不同的点：
如下图以节点2为根节点，节点1，3，6。节点1为根节点，就没有无关的点，节点4为根节点，剩下的点就全为无关的点了。

为什么要分两种呢？因为无关的点注定不会经过根节点所以组成的无根节点颜色的路数就为：n*（n-1）/2。而相关联的点就要考虑是否在根节点的同一子树上了，不同子树上的点能不能组成路，因为这样注定要经过根节点如下图：对于根节点而言2，3，4，5都是不同颜色，但是组成的不经过红色的路只有2条2-4和3-5，而不是4*3/2=6；

先 进行dfs 将每个节点当做根节点 找和他相关联的（子节点） 不相同颜色的 的路 进行统计， 并且 标记为 不可用 –》》即当 考虑和他情况时 重复统计

代码:

#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);    }}