[树] hdu6035 Colorful Tree (2017 Multi-University Training Contest - Team 1 C)

@(ACM题目)[树]

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个点的树，树上每个结点都有一个颜色，对树上的每条路径，统计其包含颜色数量。求所有路径上的这个数量的和。

为了后面叙述方便，我们将“以i结点为根节点的子树”称为sub−treei

记树上路径总数为m=n(n−1)2。考察其中一种颜色对答案的贡献，若一条路径上包含了这个颜色，那么在这条路径上，该颜色对答案贡献了“1”，否则没有贡献。正难则反，我们用路径总数m减去不包含这种颜色的路径数量，就能得到这种颜色对答案的贡献。

现计算对于某种颜色的来说，在所有的路径（ n(n−1)2条）中，不包含该颜色的路径数。若将树上所有该颜色的点从树上剔除，那么树将变一个森林。森林中的每棵树上的路径，都是不包含该颜色的路径。通过森林中每棵树包含的点数，便可得到答案。

现在通过dfs得到这个信息，参考下面的代码实现讲解：

- c[i]代表颜色为i的有哪些点

- G[i]代表与结点i相连的点

- numL[i]代表结点i的dfs序（先访问根，再访问sub−treei中其它结点的顺序）

- numR[i]代表sub−treei中，dfs序最大的那个点的dfs序

- tot[i]代表sub−treei的结点数。

通过以上定义，我们可以得出，sub−treei中结点的dfs序都在numL[i]和numR[i]之间。

分别考察每一个颜色，将该颜色的各点按dfs序排列，对于sub−treei，遍历i的孩子son，考察sub−treeson。现举例说明sub−treeson的根节点son所在的“不包含该颜色的区域”可以如何得到：



则对于sub−tree1来说，依次考察它的孩子sub−tree2、sub−tree3、sub−tree4。以sub−tree2为例，2所在的不含红色的区域为2−5区域，有tot2−tot9−tot6=8−3−3=2个结点。注意不能减去tot13，这需要利用dfs序处理，详见代码。另外虚拟增加的根节点是为了找出最上面那一部分不含该颜色的区域，当然，在这个例子中，0与当前颜色直接相连，这个区域就不存在了。

代码

#include<bits/stdc++.h>
using namespace std;
typedef long long LL;
const int maxn = 2e5+8;
int n, order;
vector<int> c[maxn], G[maxn];//c[i]: nodes with color i; e[i]: edges from/to node i
int par[maxn];//par[i]: parent of node i;
int numL[maxn], numR[maxn];
//numL[i]: minimum dfs order in sub-tree i(also the id of root); numR[i]: maximum dfs order in sub-tree i
int tot[maxn];//tot[i]: the number of nodes in sub-tree i
bitset<maxn> vis;

void dfs(int cur)
{
vis[cur] = true;
numL[cur] = ++order;
tot[cur] = 1;
for(int i = 0; i < G[cur].size(); ++i)
{
int &son = G[cur][i];
if(vis[son] == true) continue;
par[son] = cur;
dfs(son);
tot[cur] += tot[son];
}
numR[cur] = order;
}

bool cmp(int x, int y){return numL[x] < numL[y]; }

int main()
{
int Case = 1;
while(~scanf("%d", &n))
{
vis.reset();
order = -1;
for(int i = 0; i <= n; ++i)
{
c[i].clear();
G[i].clear();
}

int x, y;
for(int i = 1; i <= n; ++i)
{
scanf("%d", &x);
c[x].push_back(i);
}

G[0].push_back(1);
for(int i = 1; i < n; ++i)
{
scanf("%d%d", &x, &y);
G[x].push_back(y);
G[y].push_back(x);
}

par[0] = -1;
dfs(0);
//        for(int i = 1; i <= n; ++i) cout << "order" << i << ": " << numL[i] << ' ' << numR[i] << endl;
LL res = 0;
for(int i = 1; i <= n; ++i)//color
{
if(c[i].empty()) continue;
//            cout << "color: " << i << endl;
res += (LL)n*(n-1)/2;
c[i].push_back(0);
sort(c[i].begin(), c[i].end(), cmp);
for(int j = 0; j < c[i].size(); ++j)
{
int cur = c[i][j];//cur node
//                cout <<  "cur node: " << cur << endl;;
for(int k = 0; k < G[cur].size(); ++k)//cur children
{
int son = G[cur][k];
if(son == par[cur]) continue;
//                    cout << "son: " << son << endl;
LL m = tot[son];
numL[n+1] = numL[son];
//                    cout << "m1: " << m << endl;
vector<int>::iterator it = c[i].begin();
while(true)
{
it = lower_bound(it, c[i].end(), n+1, cmp);
if(it == c[i].end() || numL[*it] > numR[son]) break;
m -= tot[*it];
numL[n+1] = numR[*it] + 1;
}
//                    cout << "m2: " << m << endl;
res -= m*(m-1)/2;
}
}
}
printf("Case #%d: %lld\n", Case++, res);
}
return 0;
}