Educational Codeforces Round 129 (Rated for Div. 2) F. Unique Occurrences
链接
题意:给出一颗树,每两个点之间的价值是两个点之间权值的个数,记为
f
(
v
,
u
)
f(v,u)
f(v,u),规定
v
<
u
v<u
v<u,求
∑
f
(
v
,
u
)
\sum f(v,u)
∑f(v,u)
思路:可以看出,对于每条边,他的贡献就是把它去掉之后左右两个联通块大小的乘积,现在我们就要想办法把这种联通块的大小统计出来,对于每条边,我们记录他的权值,上次出现在哪个地方,维护
u
−
f
a
[
u
]
u-fa[u]
u−fa[u]权值相同的点的个数,初始化f,就可以直接处理好这个边权是第一次出现的情况,然后后面的情况直接从siz减就可以得到
#include<bits/stdc++.h>
using namespace std;
#define int long long
const int N = 5e5+10;
int head[N], idx;
struct Edge{int to, nxt, d;}e[N << 1];
void add(int u, int v, int d) {e[++idx].to = v, e[idx].nxt = head[u], e[idx].d = d, head[u] = idx;}
int f[N << 1], loc[N << 1], lst[N << 1], siz[N], n;
void dfs(int u, int fa)
{
siz[u] = 1;
for (int i = head[u]; i; i = e[i].nxt) {
int v = e[i].to, d = e[i].d;
if (v == fa) continue;
lst[v] = loc[d]; loc[d] = v;
dfs(v, u);
loc[d] = lst[v]; siz[u] += siz[v];
}f[u] += siz[u]; f[lst[u]] -= siz[u];
}
signed main()
{
cin >> n;
for (int i = 1; i < n; i++) {
int u, v, d;
cin >> u >> v >> d;
add(u, v, d); add(v, u, d);
}
for (int i = 1; i <= n; i++) {
f[i + n] = n;
loc[i] = i + n;
}
dfs(1, -1);
int ans = 0;
for (int i = 2; i <= n; i++) {
ans += (f[i] * f[lst[i]]);
}cout << ans << endl;
return 0;
}