题意:
给你一棵树, 树上的路权为 这条路上不同颜色的个数。 求所有路径的权值之和。
思路:
感觉好难。。
其实这个题一遍dfs 就可以办了。
可以反过来思考,对于每一种颜色求出 有多少条路 没有经过这个颜色。 最后用总路数减去即可。
令sum[i] 表示 以i 颜色为根的子树总大小。
siz[u]表示以u 为根的子树的大小。
怎么找呢?
假如我们枚举到u结点了。 对于u 的每个孩子, 我们找出他们的sum[col[u] ] 中间的位置差就是一个连通分块, 他们组成的路 都不包含颜色u。
这样把所有孩子都累加起来, 在更新sum[col[u] ]。
但是这样统计 会算出整体sum[col[u] ] 并不是这条链的。
这里的pre 变量就比较巧妙了。
在dfs 这个点之前,我们令pre = sum[ col[u] ] 就是 这个点之前的所有sum[col[u]]
然后更新完u的孩子后的sum[col[u] ] - pre 就是这条链的。
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <vector>
using namespace std;
const int maxn = 200000 + 10;
vector<int>g[maxn];
int col[maxn];
long long sum[maxn];
int siz[maxn];
int vis[maxn];
int num;
long long ans;
void dfs(int cur, int fa){
siz[cur] = 1;
int pre = sum[col[cur] ];
bool ok = 0;
for (int i = 0; i < g[cur].size(); ++i){
int v = g[cur][i];
if (v == fa) continue;
ok = 1;
dfs(v, cur);
siz[cur] += siz[v];
long long tmp = siz[v] - (sum[col[cur] ] - pre);
ans += (long long)tmp*(long long)(tmp-1) / 2;
sum[col[cur] ] += tmp;
pre = sum[col[cur] ];
}
sum[col[cur] ] += 1;
}
int main(){
int n, ks = 0;
while(~scanf("%d",&n)){
memset(vis,0,sizeof vis);
num = ans = 0;
for (int i = 1; i <= n; ++i){
scanf("%d", &col[i]);
if (!vis[col[i] ]){
vis[col[i] ] = 1;
++num;
}
sum[i] = 0;
g[i].clear();
}
for (int i = 1; i < n; ++i){
int x,y;
scanf("%d %d",&x, &y);
g[x].push_back(y);
g[y].push_back(x);
}
ans = 0;
dfs(1, -1);
for (int i = 1; i <= n; ++i){
if (vis[i]){
long long tmp = siz[1] - sum[i];
ans += tmp * (tmp-1) / 2;
}
}
ans = (long long)num * (long long)n * (long long)(n-1) / 2 - ans;
printf("Case #%d: %I64d\n", ++ks, ans);
}
return 0;
}
|
Colorful TreeTime Limit: 6000/3000 MS (Java/Others) Memory Limit: 131072/131072 K (Java/Others)Total Submission(s): 1733 Accepted Submission(s): 713
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
Sample Output
Source
Recommend
|