[CF1156D] 0-1-Tree - 并查集
Description
给定 n 个点的边权为 0 或 1 的树,一条合法的 x->y 路径满足一旦经过 1 边就不能再经过 0 边(换言之边权必须单调不降),求满足条件的路径数量。
Solution
所有答案点对有三种情况:全 0,全 1,一些 0 + 一些 1
前两种可以用并查集直接处理掉,假设这个 0 连通块中的点数为 x,那么贡献为 x(x-1),1 同理
对于第三种,枚举从 0 转 1 的点 i,假设这个 0 连通块中的点数为 x,假设这个 1 连通块中的点数为 y,则贡献为 (x-1)(y-1)
顺便说一下 dp 思路的基本想法:设 (f[i][0/1/2/3]) 表示自底(未必是叶子)到 i 的各种路径的条数,计算出 f 后,根据 f 讨论合并方式计算答案(似乎更容易想到就是细节颇多)
#include <bits/stdc++.h>
using namespace std;
#define int long long
const int N = 1000005;
struct DSU
{
int fa[N], sz[N];
int find(int p)
{
return p == fa[p] ? p : fa[p] = find(fa[p]);
}
void merge(int p, int q)
{
p = find(p);
q = find(q);
if (p != q)
fa[p] = q;
}
void solve(int n)
{
for (int i = 1; i <= n; i++)
sz[find(i)]++;
}
int size(int x)
{
return sz[find(x)];
}
} dsu[2];
int n;
signed main()
{
ios::sync_with_stdio(false);
cin >> n;
for (int i = 1; i <= n; i++)
dsu[0].fa[i] = dsu[1].fa[i] = i;
for (int i = 1; i < n; i++)
{
int u, v, w;
cin >> u >> v >> w;
dsu[w].merge(u, v);
}
int ans = 0;
dsu[0].solve(n);
dsu[1].solve(n);
for (int i = 1; i <= n; i++)
{
ans += dsu[0].size(i) * dsu[1].size(i) - 1;
}
cout << ans << endl;
}