2017 Multi-University Training Contest - Team 1 HDU 6035 Colorful Tree

题目来戳呀

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个节点，每个节点涂一个颜色，节点间部分有路径形成树，求所有路径上所经过节点的不同颜色数目总和。

想法：

正向考虑时，要考虑每条路径上存在的所有颜色数目，路径多数目多，所以放弃正向思维，选择逆向思维。

所有路径的颜色总和⇌通过某颜色的路径数总和
通过某颜色的路径数总和=总路径数-不经过此颜色的路径数

→ 不经过此颜色的路径数。
不经过此颜色的路径数⇌树中删去此颜色构成的连通块大小
不经过此颜色的路径数=在除去颜色形成连通块之外的路径数+连通块内形成的路径数
→ 树中删去此颜色构成的连通块大小
连通块大小=子树所有节点-以颜色为根的子树节点个数

注：代码dfs函数已经求出连通块内的路径数，在主函数最后求的是连通块之外的路径数。

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<vector>
#include<cmath>
#include<iostream>
using namespace std;
const int maxn=2e5+10;
vector<int>edge[maxn];//与结点i相连的结点 边的数组
int siz[maxn];//以节点i为根的子树节点个数
int vis[maxn];
int col[maxn];
int sum[maxn];//以颜色i为根的子树且不包含此颜色的总和
typedef long long ll;
ll ans,n;
void dfs(int u,int fa)//（子节点，父节点）
{
    siz[u]=1;
    int pre=sum[col[u]];
    int add=0;
    for(int i=0;i<edge[u].size();++i)
    {
        int v=edge[u][i];//edge[u][i]表示与u相连的第i个子节点,v是整棵树中的节点序号
        if(v==fa)//找子节点
        {//如果是父节点 就继续
           continue;
        }
        else
        {
            dfs(v,u);
        }

        siz[u]+=siz[v];

        ll sonu=sum[col[u]]-pre; //在子树v中，以颜色col[u]为根的最高的子树的大小
        add+=sonu;//add记录当前所有子树中以col[u]为根的节点总数 
        ll tmp=siz[v]-sonu;//子树v中，除了u颜色之外的节点个数
        ans-=tmp*(tmp-1)/2; //表示v子树中，各个节点互相到达，无需经过col[u]
        pre=sum[col[u]];//表示记录上一次的sum
    }

    sum[col[u]]+=siz[u]-add;//sum[col[u]]加上以u为根，属于这个节点的连通块的节点总数
}
int main()
{
    int cas=1;
    while(~scanf("%lld",&n))
    {
        memset(sum,0,sizeof sum);
        memset(col,0,sizeof col);
        memset(siz,0,sizeof siz);
        memset(vis,0,sizeof vis);
        int cnt=0;//颜色总种数
        for(int i=1;i<=n;++i)
        {
            scanf("%d",&col[i]);
            if(!vis[col[i]])
            {
                vis[col[i]]=1;
                cnt++;
            }
        }
         int u,v;
         for(int i=1;i<=n;++i)
         {
             edge[i].clear();
         }
         for(int i=1;i<n;++i)
          {
              scanf("%d %d",&u,&v);
              edge[u].push_back(v);//构成u v 之间的无向边
              edge[v].push_back(u);
          }

          ans=cnt*(n)*(n-1)/2;//所有颜色都存在的总路径数
         dfs(1,0);
        for(int i=1;i<=n;++i)
        {
            if(!sum[i])
            {
                continue;
            }
            ll res=n-sum[i];
            ans-=(res-1)*res/2;//此步骤减去的是 在i颜色节点的 上面部分，有n-sum[i]个节点，这些节点互通不需要经过i颜色节点。
        }
        printf("Case #%d: %lld\n",cas++,ans);
    }
}

而标程其实是一样逆向的思想，只是在求子树中不含此颜色的节点数采取的是用dfs序来判断的方式。
dfs序：将树按dfs的方式排成一个序列，此点在序列中的序号称为dfs序。

①找出颜色中dfs序大于等于子根dfs序的位置
②若该位置指向该颜色的最后一个位置，代表已经是该子树的最后
③若该位置大于该子根的最大的dfs序，代表已经不属于这棵子树
若不符合②③这两种情况，就删去这个位置，因为这个位置含有以此位置为根的子树。

注：标程是在c++11才能运行的。

#include <bits/stdc++.h>    
using namespace std;    
typedef long long LL;    
const int N = 200005;    

int n , ca;    
vector<int> e[N] , c[N];// e[i]与结点i相连的结点(边数组), c[i]涂有颜色i有哪些点
int L[N] , R[N] , s[N] , f[N];    //L[i]结点i的dfs序,  R[i]以i为根结点的子树中dfs序最大的点的dfs序, S[i]以i为根节点的树共有多少个结点, f[i]结点i的父亲结点

void dfs(int x , int fa , int &&ncnt)    
{    
    L[x] = ++ ncnt;    
    s[x] = 1 , f[x] = fa;    
    for (auto &y : e[x])    
    {    
        if (y != fa)    
        {    
            dfs(y , x , move(ncnt));    
            s[x] += s[y];    
        }    
    }    
    R[x] = ncnt;    
}    

bool cmp(const int& x , const int& y)    
{    
    return L[x] < L[y];    
}    

void work()    
{    
    for (int i = 0 ; i <= n ; ++ i)   //从0开始 加了一个虚根 表示子树路径中不存在的颜色== 
    {    
        c[i].clear();    
        e[i].clear();    
    }    
    for (int i = 1 ; i <= n ; ++ i)    
    {    
        int x;    
        scanf("%d" , &x);    
        c[x].push_back(i);    
    }    
    for (int i = 1 ; i < n ; ++ i)    
    {    
        int x , y;    
        scanf("%d%d" , &x , &y);    
        e[x].push_back(y);    
        e[y].push_back(x);    
    }    
    e[0].push_back(1);    
    dfs(0 , 0 , 0);    
    LL res = (LL)n * n * (n - 1) / 2; //假设n个结点颜色均不同,每种颜色都贡献了所有的路径    
    for (int i = 1 ; i <= n ; ++ i)    
    {    
        if (c[i].empty()) //如果为空,减一次(没有这种颜色,减掉这种颜色所贡献的路径(全部))    
        {    
            res -= (LL)n * (n - 1) / 2;    
            continue;    
        }    
        c[i].push_back(0);    
        sort(c[i].begin() , c[i].end() , cmp);    
        for (auto &x : c[i])    
        {    
            for (auto &y : e[x])代表遍历e[x]中所有的元素，每次遍历时都把元素成为y    
            {    
                if (y == f[x]) //若为父亲结点则continue(只考虑孩子结点)    
                    continue;    
                int size = s[y];    
                int k = L[y];    
                while (1)    
                {    
                    L[n + 1] = k;    
                    auto it = lower_bound(c[i].begin() , c[i].end() , n + 1 , cmp); //找以孩子结点为根节点的子树中,相同颜色的结点    
                    if (it == c[i].end() || L[*it] > R[y])    
                    {    
                        break;    
                    }    
                    size -= s[*it];    
                    k = R[*it] + 1; //另一棵子树    
                }    
                res -= (LL)size * (size - 1) / 2;    
            }    
        }    

    }    
    printf("Case #%d: %lld\n" , ++ ca , res);    
}    

int main()    
{    
    while (~scanf("%d" , &n))    
    {    
        work();    
    }    
    return 0;    
}  

ps：哇这个题真的是看了好久啊QAQ
看懂了才发现没有那么难，是自己不够安静去研究每一个变量的意思，其实自己静下心来跑一遍就好了啊0.0