2017多校训练赛第一场 HDU 6035 Colorful Tree （dfs+正序统计）

Colorful Tree

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

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 nodei 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

Source

2017 Multi-University Training Contest - Team 1

        这题，好像是没做出的所有题目中，最容易做出来的题目，但是当时真的静不下心啊……

        统计所有路径上的颜色数量，直接枚举所有的路径条数显然不行，比较惯用的计算方法是统计每一个点对最后结果的贡献。这里题解是用了补集的思想，反过来求，我给出一种另外的思路，不反着做，直接正着统计和。

        首先，我们规定，对于同一条路径上的两个相同的颜色，我们只统计dfs序靠前的那个。然后，对于一个点i，我们可以分成两种情况讨论。第一种情况，我们只计算i的子树范围内的方案数，这个还是比较容易计算的，其中一部分等于sigma(sigma(sz[j]*sz[k]))(1<j<k<i的儿子个数)，其中sz[i]表示以i为根的子树的大小，两个sigma分别对j和k求和，即任意两个儿子的子树节点数相乘，表示所有从i的一个儿子到i的另一个儿子的路径；另一部分为以i为根的子树个数减1，表示从i出发到i的任意一个后代的路径。

        以上是第一种情况，由于他相对于它的所有后代dfs序列都靠前，所以不用考虑颜色相同的情况。第二种情况则考虑从以i为根的子树出去连到子树外任意一条边的方案数。显然，当不考虑相同颜色的时候，结果就是(n-sz[i])*sz[i]，但是要考虑有颜色相同的情况。我们假设i的相同颜色的第一个祖先为fa，则我们对应每种颜色开一个栈st，st的初始值为fa的儿子f的大小sz[f]，这个f为i的祖先。这样我们计算的时候就不是用(n-sz[i])*sz[i]而是用(st.top()-sz[i])*sz[i]，就可以保证进过fa的路径不被计算进去。但是这样还有缺陷，因为即使fa不被算入，f可能有其他儿子颜色与i相同并且先被计算过。这种情况，我们就考虑每次计算完i之后，对栈顶的元素减去sz[i]，这样栈顶的数字就可以保证没有颜色相同的情况。

        最后还有一种特殊的情况，如果i没有与其颜色相同的祖先呢？同样的，我们用类似的方法，(rt[col[i]]-sz[i])*sz[i]，rt类似与之前的栈st，一样也是每次计算完i之后减去sz[i]即可。后来看了一下官方题解，发现其实每种颜色的栈首先加进去一个虚拟的祖先就可以解决这种问题了。具体见代码：

#include<iostream>#include<cstring>#include<cstdio>#include<vector>#include<stack>#define LL long long#define N 200100using namespace std;stack<int> st[N];vector<int> g[N];int col[N],sz[N],rt[N],n,root;LL ans;void getsize(int x,int fa){    sz[x]=1;    for(int i=0;i<g[x].size();i++)    {        int y=g[x][i];        if (y!=fa)        {            getsize(y,x);            sz[x]+=sz[y];        }    }}void dfs(int x,int fa){    int sum=0;    for(int i=0;i<g[x].size();i++)    {        int y=g[x][i];        if (y!=fa)        {            st[col[x]].push(sz[y]);            dfs(y,x);            ans+=(LL)sum*sz[y];//以儿子为根的子树大小两两相乘            sum+=sz[y];            st[col[x]].pop();        }    }    ans+=(LL)sum;//还要加上所有后代个数    if (!st[col[x]].empty())    {        st[col[x]].top()-=sz[x];        ans+=(LL)st[col[x]].top()*sz[x];//第二部分，子树内部到外部的贡献    } else ans+=(LL)(sz[root]-sz[x]-rt[col[x]])*sz[x],rt[col[x]]+=sz[x];//没有相同颜色祖先的情况}int main(){    int T_T=0;    while(~scanf("%d",&n))    {        memset(rt,0,sizeof(rt));        memset(g,0,sizeof(g));        for(int i=1;i<=n;i++)            scanf("%d",&col[i]);        for(int i=1;i<n;i++)        {            int u,v;            scanf("%d%d",&u,&v);            g[u].push_back(v);            g[v].push_back(u);        }        ans=0;        root=1;        getsize(1,0);        dfs(1,0);        printf("Case #%d: %I64d\n",++T_T,ans);    }    return 0;}