Colorful Tree
Time Limit: 6000/3000 MS (Java/Others) Memory Limit: 131072/131072 K (Java/Others)Total Submission(s): 973 Accepted Submission(s): 375
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.
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.
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
Source
解题思路:题意可以转化为所有颜色种类*所有路径数量-每一种颜色有多少路径没有经过,然后dfs就可以O(n)扫出每一种颜色的贡献, Size[k]表示以节点k为根节点的子树的大小,sum[i]表示以颜色i为根节点的子树的大小和,具体看代码
#include <iostream>
#include <cstdio>
#include <cstring>
#include <string>
#include <algorithm>
#include <map>
#include <cmath>
#include <set>
#include <stack>
#include <queue>
#include <vector>
#include <bitset>
#include <functional>
using namespace std;
#define LL long long
const int INF = 0x3f3f3f3f;
const LL mod=1e9+7;
int n;
int a[200009],sum[200009];
int Size[200009],vis[200009],x,y;
int s[200009],nt[400009],e[400009];
LL ans;
void dfs(int k,int fa)
{
Size[k]=1;
int sum1=0;
for(int i=s[k];~i;i=nt[i])
{
if(e[i]==fa) continue;
int p=sum[a[k]];
dfs(e[i],k);
Size[k]+=Size[e[i]];
sum1+=sum[a[k]]-p;
int pp=Size[e[i]]-(sum[a[k]]-p);
ans-=1LL*(pp-1)*pp/2;
}
sum[a[k]]+=Size[k]-sum1;
}
int main()
{
int cas=0;
while(~scanf("%d",&n))
{
memset(sum,0,sizeof sum);
memset(s,-1,sizeof s);
memset(vis,0,sizeof vis);
int cnt=0,tot=0;
for(int i=1;i<=n;i++)
{
scanf("%d",&a[i]);
if(!vis[a[i]]) tot++;
vis[a[i]]=1;
}
for(int i=1;i<n;i++)
{
scanf("%d%d",&x,&y);
nt[cnt]=s[x],s[x]=cnt,e[cnt++]=y;
nt[cnt]=s[y],s[y]=cnt,e[cnt++]=x;
}
ans=1LL*n*(n-1)/2*tot;
dfs(1,0);
for(int i=1;i<=n;i++)
{
if(!vis[i]) continue;
ans-=1LL*(n-sum[i])*(n-sum[i]-1)/2;
}
printf("Case #%d: %lld\n",++cas,ans);
}
return 0;
}