17暑假多校联赛1.3 HDU 6035 Colorful Tree

Colorful Tree

Time Limit: 6000/3000 MS (Java/Others)
Memory Limit: 131072/131072 K (Java/Others)
 

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

31 2 11 22 361 2 1 3 2 11 21 32 42 53 6

Sample Output

Case #1: 6Case #2: 29

 
题目网址：http://acm.hdu.edu.cn/showproblem.php?pid=6035
 

分析

题意：给一颗有n个节点的树，每个节点有一种颜色，颜色用整数来表示，定义树上路径u到v的价值为路径上不同颜色的节点的数量，问所有路径的总价值为多少
问题可以转变为计算每种颜色有多少条路径经过，又可以转变为所有情况减去每种颜色没有经过的路径数
详解见代码注释
C++中的1LL
 

代码

#include <bits/stdc++.h>using namespace std;const int N = 4e5+100;typedef long long ll;int c[N],vis[N];///c存放各个节点的颜色，vis来标记各种颜色是否出现过vector<int> e[N];///e存放各条路径ll sum[N],size[N];///sum存放已经查询过的路径(包含此点)颜色为i的节点数///整个遍历之后，sum存放的是以i为根节点的子节点数///size存放各个点为根节点时的子节点数目ll ans;///ans存放每种颜色没有经过的路径数void dfs(int x,int y){    size[x]=1;    sum[c[x]]++;    ll pre=sum[c[x]];    ///pre为已经查询过的路径(不包含此点)颜色为i的节点数    for(int i=0; i<e[x].size(); i++)    {        if(e[x][i]==y)///如果等与他的父节点就不用遍历            continue;        dfs(e[x][i],x);///遍历x的子节点        size[x]+=size[e[x][i]];        ///x的子节点包含x的子节点的子节点        ll count=size[e[x][i]]-(sum[c[x]]-pre);        ///count存放的是以e[x][i]为根结点的颜色不为c[x]的点的个数        ans=ans+(1LL*count*(count-1))/2;        ///ans加上count这些点组成的路径数        sum[c[x]]+=count;        ///加上这些点，避免计算另一个子结点时结果重复        pre=sum[c[x]];///pre也随sum变化，避免重复计算    }}int main(){    int n,cas=1;    while(~scanf("%d",&n))    {        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]])            {                vis[c[i]]=1;///颜色c[i]出现过，标记为1                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);            ///将路径存入e中        }        dfs(1,0);///从根节点开始遍历        ll ANS = 1LL*num*((1LL)*n*(n-1))/2;        ///ANS计算的是所有情况        ///颜色数乘以路径数，各种颜色在每条路径都存在的数目        for(int i=1; i<=n; i++)        {            if(vis[i])            {                ll ct=n-sum[i];///ct为不以颜色i为根结点的节点个数                ans+=ct*(ct-1)/2;///ans加上这些点构成的路径            }        }        printf("Case #%d: %lld\n", cas++, ANS-ans);        ///所有情况减去每种颜色没有经过的路径数    }}