【多校训练】hdu 6035 Colorful Tree 树状数组

Colorful Tree

Time Limit: 6000/3000 MS (Java/Others)    Memory Limit: 131072/131072 K (Java/Others)
Total Submission(s): 2100    Accepted Submission(s): 892

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-1)/2)，再减去实际上多统计的就为答案

-> 怎么求解多统计的贡献呢？如果对于一种颜色，其相邻的点中都没有这个颜色，那么这些相邻点构成的连通块中的路径都为多统计的贡献

->需要用树状数组来计算多余的贡献，代码中还用到了很多技巧，具体实现见代码。

//
//  main.cpp
//  1003
//
//  Created by zc on 2017/8/2.
//  Copyright © 2017年 zc. All rights reserved.
//

#include <iostream>
#include<cstdio>
#include<cstring>
#include<cmath>
#include<algorithm>
#include<vector>
#define ll long long
using namespace std;
const int N=550000;
vector<int>r[N];
int n,sz[N],v[N];
ll c[N],ans,sum[N];

ll dfs(int t,int fa)
{
    sz[t]=1;
    for(int i=0;i<r[t].size();i++)
    {
        int u=r[t][i];
        if(u==fa)   continue;
        ll last=sum[c[t]];//保存没计算没计算子树之前，以c[t]为根的子树的总大小
        sz[t]+=dfs(u,t);
        ll add=sum[c[t]]-last;//add=计算子树后的sum-没计算前的sum=该子树拥有的以c[t]颜色为根的子树的总大小啊
        ans+=(sz[u]-add)*(sz[u]-add-1ll)/2ll;//sz[u]-add=该子树中没有c[t]颜色的连通块大小
        sum[c[t]]+=sz[u]-add;//sum加上连通块
    }
    sum[c[t]]++;//加上自己
    return sz[t];
}

int main(int argc, const char * argv[]) {
    int kase=0;
    while(~scanf("%d",&n))
    {
        int m=0;
        memset(v,0,sizeof(v));
        for(int i=1;i<=n;i++)
        {
            scanf("%lld",&c[i]);
            r[i].clear();
            if(v[c[i]]==0)
                v[c[i]]=1,m++;
        }
        for(int i=0;i<n-1;i++)
        {
            int u1,u2;
            scanf("%d%d",&u1,&u2);
            r[u1].push_back(u2);
            r[u2].push_back(u1);
        }
        ans=0;
        memset(sum,0,sizeof(sum));
        dfs(1,-1);
        for(int i=1;i<=n;i++)//需要加上整棵树作为子树统计的各种颜色的连通块
            if(v[i])    ans+=(n-sum[i])*(n-sum[i]-1ll)/2ll;
        printf("Case #%d: %lld\n",++kase,(ll)m*n*(n-1ll)/2ll-ans);
    }
}