题目传送门: http://acm.hdu.edu.cn/showproblem.php?pid=6035
题意:
给出一颗树, 树上路径权值等于路径上不同颜色的数量, 求所有路径权值之和;
解题思路:
大佬讲的很清楚了:
AC code:
- #include <bits/stdc++.h>
- #define INF 0x3f3f3f3f
- #define LL long long
- using namespace std;
- const int MAXN = 2e5+10;
- int c[MAXN];
- bool vis[MAXN];
- vector<int>e[MAXN];
- LL sum[MAXN], Size[MAXN];
- LL ans;
- void dfs(int x, int fa)
- {
- Size[x] = 1;
- sum[c[x]]++;
- LL pre = sum[c[x]];
- int len = e[x].size();
- for(int i = 0; i < len; i++){
- if(e[x][i] == fa) continue;
- dfs(e[x][i], x);
- Size[x]+=Size[e[x][i]];
- LL Count = Size[e[x][i]]-(sum[c[x]]-pre);
- ans = ans+(Count*(Count-1))/2;
- sum[c[x]]+=Count;
- pre = sum[c[x]];
- }
- }
- int main()
- {
- int N, cas = 1;
- LL num = 0, ct = 0;
- while(~scanf("%d", &N)){
- num = 0; ans = 0;
- memset(sum, 0, sizeof(sum));
- memset(vis, 0, sizeof(vis));
- for(int i = 1; i <= N; i++){
- e[i].clear();
- scanf("%d", &c[i]);
- if(!vis[c[i]]){
- vis[c[i]] = true;
- num++;
- }
- }
- int u, v;
- for(int i = 1; i < N; i++){
- scanf("%d %d", &u, &v);
- e[u].push_back(v);
- e[v].push_back(u);
- }
- dfs(1, -1);
- LL res = (LL)(1LL*N*(N-1))/2*num;
- for(int i = 1; i <= N; i++){
- if(vis[i]){
- ct = N-sum[i];
- ans+=ct*(ct-1)/2;
- }
- }
- printf("Case #%d: %lld\n", cas++, res-ans);
- }
- return 0;
- }
- View Code
来源: http://www.bubuko.com/infodetail-3044908.html