[CF1156D] 0-1-Tree

[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;
}