（多校） hdu 6035 Colorful Tree

Colorful Tree

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

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*n,再减去非法的情况，也就是每种颜色相同的连通块的情况，就是最后答案。
做法就是先求出这棵树的dfs序，然后用线段树去维护每个点有多少个点与它连通，之后用总答案的减去这种颜色的非法情况，对每种颜色都做一次这样的操作，就是最后的答案。

#include <stdio.h>
#include <iostream>
#include <algorithm>
#include <string>
#include <string.h>
#include <set>
#include <map>
#include <queue>
#include <stack>
#include <math.h>
#include <vector>
#include <windows.h>
using namespace std;
const int maxn = 200005;
long long n;
int fr,tim;
int col[maxn],vis[maxn],head[maxn];
int l[maxn],r[maxn],fa[maxn];
vector<int>op[maxn];
struct pp
{
    int node,next;
}E[maxn*4];
struct p
{
    int l,r,lz,val;
}tree[maxn*4];
void Insert(int x,int y)
{
    E[fr].node=y;
    E[fr].next=head[x];
    head[x]=fr;
    fr++;
}
void build(int i,int l,int r)
{
    tree[i].l=l;tree[i].r=r;tree[i].val=0;
    if (l==r) return;
    int mid=(l+r)/2;
    build(i<<1,l,mid);
    build((i<<1)+1,mid+1,r);
}
void update(int i,int l,int r,int v)
{

    if (l==tree[i].l&&tree[i].r==r)
    {
        tree[i].val+=v;
        return;
    }
    if (tree[i].lz!=-1){
      tree[i<<1].lz=tree[i].lz;tree[(i<<1)+1].lz=tree[i].lz;
      tree[i*2].val=0;tree[i*2+1].val=0;
      tree[i].lz=-1;
      tree[i].val=0;
    }
    int mid=(tree[i].l+tree[i].r)/2;
    if (r<=mid) {update(i*2,l,r,v);}
    else if (l>mid) {update(i*2+1,l,r,v);}
    else {update(i*2,l,mid,v);update(i*2+1,mid+1,r,v);}
    tree[i].val=tree[i*2].val+tree[i*2+1].val;
}
int find(int i,int l,int r)
{
    if (tree[i].l==l&&tree[i].r==r)
    {
        return tree[i].val;
    }
    if (tree[i].lz!=-1)
    {
        tree[i*2].lz=tree[i].lz;tree[i*2+1].lz=tree[i].lz;
        tree[i*2].val=0;tree[i*2+1].val=0;
        tree[i].lz=-1;
        tree[i].val=0;
    }
    int mid=(tree[i].l+tree[i].r)/2;
    if (r<=mid) {return find(i*2,l,r);}
    else if(l>mid) {return find(i*2+1,l,r);}
    else return find(i*2,l,mid)+find(i*2+1,mid+1,r);
}
void dfs(int root)
{
    vis[root]=1;
    l[root]=++tim;
    for (int i=head[root];i!=-1;i=E[i].next)
    {
        int nod=E[i].node;
        if (!vis[nod])
        {
            fa[nod]=root;
            dfs(nod);
        }
    }
    r[root]=tim;
    op[col[root]].push_back(root);

}
int main()
{
    int T=0;
    while (scanf("%I64d",&n)!=EOF)
    {
        T++;
        fr=0;
        tim=0;
        memset(head,-1,sizeof(head));
        memset(vis,0,sizeof(vis));
        for (int i=0;i<maxn;i++) op[i].clear();
        for (int i=1;i<=n;i++) scanf("%d",&col[i]);
        for (int i=1;i<n;i++)
        {
            int x,y;
            scanf("%d%d",&x,&y);
            Insert(x,y);
            Insert(y,x);
        }
        fa[1]=0;
        dfs(1);
        build(1,1,n);

        long long ans=0;
        for (int i=1;i<=n;i++)
        {
            if (op[i].size()>0)
            {
                tree[1].lz=0;
                long long tmp=n*(n-1)/2;
                for (int j=0;j<op[i].size();j++)
                {
                    int x=op[i][j];
                    update(1,l[x],l[x],1);
                    for (int k=head[x];k!=-1;k=E[k].next)
                    {
                        int nod=E[k].node;
                        if (nod==fa[x]) continue;
                        long long  g=find(1,l[nod],r[nod]);
                        g=r[nod]-l[nod]+1-g;
                        tmp-=g*(g-1)/2;
                        update(1,l[nod],l[nod],g);
                    }
                }
                if (op[i][op[i].size()-1]!=1)
                {
                    long long g=find(1,l[1],r[1]);
                    g=r[1]-l[1]+1-g;
                    tmp-=g*(g-1)/2;
                }
                    ans+=tmp;

            }
        }
        printf("Case #%d: %I64d\n",T,ans);
    }
}