(本文转载,原帖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;
}