2017多校1 1003 HDU 6035 树形dp

（本文转载，原帖http://blog.youkuaiyun.com/Bahuia/article/details/76141574）

题意：

题目链接：http://acm.hdu.edu.cn/showproblem.php?pid=6035 

一棵n个结点的树，每个结点都有颜色，定义两点之间的路径长度为路径上出现的不同颜色数目，求树上所有路径的长度和。

思路：

“真的难”系列。 
首先这题肯定是算贡献，也就是计算出每种颜色参与了多少条路径，但这样正面考虑并不容易，不妨从反面考虑，计算每种颜色没有参与多少路径，然后拿 (路径总数 * 颜色总数) - 没参与的贡献，就是答案了。 
对于一种颜色x，怎么计算没参与的路径数目呢，很显然，对于每个不包含颜色x的连通块中任意两点路径都是x不参与的贡献，那么问题就转化为，对于任意一种颜色x，需要求出每个不包含x的连通块大小。 
这里利用了树形dp，对于结点u，若u的颜色为x，那么dfs(u)的过程中，我们就想知道对于u的每个儿子v构成的子树中最高的一批颜色也为x的结点是哪些，要是知道这些结点子树的大小，只要拿子树v的大小减去这些节点子树的大小，就可以得到包括v在内的没有颜色x的连通块大小。 
举个例子，如图： 
 
对于结点1，我们想求出与1相邻的且不是黑色的结点组成的连通块大小，此时就要dfs结点1的两个儿子2和3，若我们处理出来结点2中，最高的一批黑色结点（4，8）的所构成子树的大小分别为2和1，那么我们拿2为根子树的大小5，减去2，减去1，就得到了结点2不包含黑色结点的连通块大小是5-2-1=2，也就是图中的2和5组成的连通块。同理对3，可以求得连通块大小是2（{3，6}）。 
计算出了子树u中连通块大小对答案的贡献之后，就要将当前最高一批的黑色结点更新为u，最高黑色节点构成的子树大小sum加上子树u中没有计算的那部分大小5（{1,2,3,5,6}）。 
具体实现，看代码，sum[x]表示的遍历到当前位置，颜色为x的高度最高一批结点为根的子树大小总和。

#include <iostream>
#include <cstdio>
#include <cctype>
#include <algorithm>
#include <cstring>
#include <string>
#include <cmath>
#include <vector>
#include <set>
#include <stack>
#include <sstream>
#include <queue>
#include <map>
#include <functional>
#include <bitset>
#include <ctime>

using namespace std;
#define pb push_back
#define mk make_pair
#define ll long long
#define ull unsigned long long
#define pii pair<int, int>
#define mk make_pair
#define fi first
#define se second
#define ALL(A) A.begin(), A.end()
#define rep(i,n) for(int (i)=0;(i)<(int)(n);(i)++)
#define repr(i, n) for(int (i)=(int)(n);(i)>=0;(i)--)
#define repab(i,a,b) for(int (i)=(int)(a);(i)<=(int)(b);(i)++)
#define reprab(i,a,b) for(int (i)=(int)(a);(i)>=(int)(b);(i)--)
#define sc(x) scanf("%d", &x)
#define pr(x) printf("x:%d\n", x)
#define fastio ios::sync_with_stdio(0), cin.tie(0)
#define frein freopen("in.txt", "r", stdin)
#define freout freopen("out.txt", "w", stdout)
#define freout1 freopen("out1.txt", "w", stdout)
//#define lb puts("")
#define lson ((rt<<1)+1)
#define rson ((rt<<1)+2)
#define mid ((l+r)/2)
#define lmid (l+(r-l)/3)
#define rmid (r-(r-l)/3)
#define debug cout<<"???"<<endl
#define pll pair<long long, long long>

const double PI = 3.1415926535897932384626433;
const ll mod = 1e9+7;
const int INF = 0x3f3f3f3f;
const double eps = 1e-12;
template<class T> T gcd(T a, T b){if(!b)return a;return gcd(b,a%b);}
const int maxn = 2e5+10;

ll cnt;
int color[maxn], sum[maxn], n, sz[maxn];
vector<int> tree[maxn];

void dfs(int u, int f)
{
    sz[u] = 1;
    ll son = 0;
    for(int i = 0; i < tree[u].size(); i++){
        int v = tree[u][i];
        if(v == f) continue;
        ll pre = sum[color[u]];
        dfs(v,u);
        sz[u] += sz[v];
        ll add = sum[color[u]]-pre;
        //printf("dfs(%d,%d): sz[%d]:%d, pre:%lld, add:%lld\n",u,f,v,sz[v],pre,add);
        add = sz[v]-add;
        cnt += add*(add-1)/2;
        son += add;
    }
    sum[color[u]] += son+1;
}

void init()
{
    memset(color, 0, sizeof(color));
    memset(sum, 0, sizeof(sum));
    for(int i = 1; i <= n; i++) tree[i].clear();
}

int main()
{
    //frein;
    //freout;
    //fastio;
    int kase = 1;
    while(scanf("%d", &n) == 1){
        init();
        for(int i = 1; i <= n; i++) scanf("%d", &color[i]);
        for(int i = 1; i < n; i++){
            int u,v; scanf("%d%d", &u, &v);
            tree[u].pb(v);
            tree[v].pb(u);
        }
        cnt = 0;
        dfs(1,1);
        int col = 0;
        for(int i = 1; i <= n; i++){
            if(sum[i] == 0) continue;
            col++;
            cnt += (n-sum[i])*(n-sum[i]-1LL)/2LL;
        }
        ll ans = n;
        ans = ans*(ans-1)/2LL*col - cnt;
        printf("Case #%d: %lld\n", kase++, ans);
    }
    return 0;
}