HDU 6035 colorful tree

题目链接

HDU 6035 colorful tree

分析

参考

首先，我们单独考虑每种颜色，对于颜色 c 　来说，他对答案的贡献就是包含这种颜色的路径条数，反过来想就是　n∗(n−1)/2−（不经过这种颜色的路径条数）,怎么求呢?我们以颜色　c 来对树分块，那么　就是用总路径条数　n∗(n−1)/2 剪掉　每一块的路径数目.不过我们只需要一次dfs就可以求出

代码

#include<bits/stdc++.h>
using namespace std;
#define pb push_back
#define mkp make_pair
#define fi first
#define se second
#define ll long long
#define M 1000000007
#define all(a) a.begin(), a.end()

const int maxn = 200100;
int n, ca, c[maxn], dfn[maxn], tot;
ll sz[maxn], ans;
vector<int> ed[maxn];
stack<int> vec[maxn];

void dfs(int t, int fa){
dfn[t] = ++tot;
sz[t] = 1;
for(auto v : ed[t]){
if(v == fa) continue;
dfs(v, t);
sz[t] += sz[v];
int tmp = sz[v];
while(!vec[c[t]].empty() && dfn[vec[c[t]].top()] > dfn[t]){
tmp -= sz[vec[c[t]].top()];
vec[c[t]].pop();
}
ans -= (ll)tmp * (tmp - 1) / 2;
}
vec[c[t]].push(t);
}

int main(){
while(~scanf("%d", &n)){
for(int i = 1; i <= n; ++i) scanf("%d", c + i);
for(int i = 1; i < n; ++i){
int u, v;
scanf("%d%d", &u, &v);
ed[u].pb(v);
ed[v].pb(u);
}
tot = 0;
ans = (ll)n * n * (n - 1) / 2;
dfs(1, 0);
for(int i = 1; i <= n; ++i){
ed[i].clear();
int tmp = n;
while(!vec[i].empty()){
tmp -= sz[vec[i].top()];
vec[i].pop();
}
ans -= (ll)tmp * (tmp - 1) / 2;
}
printf("Case #%d: %lld\n", ++ca, ans);
}
return 0;
}

dfs剪掉的是每种颜色分块后第一个顶点到叶子顶点构成的树中的块，具体来说是用颜色栈vec 记录每种颜色的访问顺序，剪掉到访问顶点u 的最近且和u 颜色相同的点,for循环减去的是每种颜色分块后第一个顶点到根上部分．

不过这份代码会超内存，因为颜色栈　vec　每次计算更新内存太大

AC code

g改进一下我们用　sum　来存每一块的累加值，用pre来记录访问　子节点　以前的变化值，再用访问后sun值来剪掉这个pre就是离他最近子块的节点数目，用总的子树节点数剪掉这个值就是当前块的节点数目.

#include<bits/stdc++.h>
#define pb push_back
#define mp make_pair
#define PI acos(-1)
#define fi first
#define se second
#define INF 0x3f3f3f3f
#define INF64 0x3f3f3f3f3f3f3f3f
#define random(a,b) ((a)+rand()%((b)-(a)+1))
#define ms(x,v) memset((x),(v),sizeof(x))
using namespace std;
const int MOD = 1e9+7;
const double eps = 1e-8;
typedef long long LL;
typedef long double DB;
typedef pair<int,int> PII;
const int maxn = 2e5+10;
const int MAX_V = 1e5+10;

int c[maxn];
std::vector<int> G[maxn];
int sz[maxn];
int sum[maxn];
LL ans;
void dfs(int u,int fa) {
sz[u] = 1;
LL pre = sum[c[u]];
LL add = 0;
for(auto v : G[u]){
if(v == fa)continue;
dfs(v,u);
sz[u] += sz[v];
LL tmp = sz[v];
// while (!vec[c[u]].empty() && dfn[vec[c[u]].top()] > dfn[v]) {
//     tmp -= sz[vec[c[u]].top()];
//     vec[c[u]].pop();
// }
// ans -= tmp*(tmp-1)/2;
tmp -= (sum[c[u]] - pre);
ans -= tmp * (tmp - 1) /2;
add += tmp;
pre  = sum[c[u]];
}
sum[c[u]] += add+1;
}

int main() {

std::ios::sync_with_stdio(false);
std::cin.tie(0);
int n;
int kase =0;
c[0] = 0;
while (cin>>n) {
for(int i=1 ; i<=n ; ++i)cin>>c[i];
for(int i=1 ; i<n ; ++i){
int u,v;cin>>u>>v;
G[u].pb(v);G[v].pb(u);
}

ans = (LL)n*n*(n-1)/2;
ms(sum,0);
dfs(1,0);
for(int i=1 ; i<=n ; ++i){
LL tmp = n - sum[i];
ans -= tmp*(tmp-1) /2;
}
for(int i=1 ; i<=n ; ++i)G[i].clear();
std::cout << "Case #" << ++kase << ": "<< ans<< '\n';

}

return 0;
}