Siyuan

「Codeforces 1156D」0-1-Tree
「Codeforces 1156D」0-1-Tree
扫描右侧二维码阅读全文
20
2019/05

「Codeforces 1156D」0-1-Tree

题目链接:Codeforces 1156D

你有一棵由 $n$ 个节点组成的树,每条边上有一个数字 $0$ 或 $1$。

假如我们从节点 $x$ 开始沿着简单路径到达节点 $y$,一旦经过了数字为 $1$ 的边就不会经过数字为 $0$ 的边。那么我们称二元组 $(x, y)$($x \ne y$)是合法的。

你的任务是找出树上合法的二元组数量。

数据范围:$2 \le n \le 2 \times 10 ^ 5$。


Solution

我们对这条路径分为 $3$ 种情况讨论,对这些情况分别计算贡献。

  1. 只包含 $0$ 边。
  2. 只包含 $1$ 边。
  3. 起初为 $0$ 边,后来为 $1$ 边。

对于 $0, 1$ 边分别建立其生成树森林,设数字为 $j$ 的第 $i$ 棵树的大小为 $size(i, j)$,那么对于三种情况分别有:

  1. 贡献为 $\sum A_{size(i, 0)}^{2}$(每个集合产生一次贡献)。
  2. 贡献为 $\sum A_{size(i, 1)}^{2}$(每个集合产生一次贡献)。
  3. 考虑枚举交界的那个点,那么在其所在 $0, 1$ 边生成树中各选一个点,贡献为 $(size(belong(i), 0) - 1) \cdot (size(belong(i), 1) - 1)$。其中 $belong(i)$ 表示点 $i$ 所在集合,当然因数字而异(每个点产生一次贡献)。

当然也可以通过树形 $\text{DP}$ 在 $\mathcal O(n)$ 的更优复杂度内求解。但是鉴于码量和细节问题,笔者选择并查集的解法。

时间复杂度:$\mathcal O(n \log n)$。


Code

#include <cstdio>

const int N = 2e5 + 5;

int n, u[N], v[N], w[N], fa[N][2], sz[N][2];

int find(int x, int c) {
    return fa[x][c] == x ? x : fa[x][c] = find(fa[x][c], c);
}
void merge(int x, int y, int c) {
    x = find(x, c), y = find(y, c);
    if (x == y) return;
    fa[y][c] = x, sz[x][c] += sz[y][c];
}
int main() {
    scanf("%d", &n);
    for (int i = 1; i <= n; i++) {
        fa[i][0] = fa[i][1] = i;
        sz[i][0] = sz[i][1] = 1;
    }
    for (int i = 1; i < n; i++) {
        int u, v, w;
        scanf("%d%d%d", &u, &v, &w);
        merge(u, v, w);
    }
    long long ans = 0;
    for (int i = 1; i <= n; i++) {
        if (fa[i][0] == i) ans += 1LL * sz[i][0] * (sz[i][0] - 1);
        if (fa[i][1] == i) ans += 1LL * sz[i][1] * (sz[i][1] - 1);
        ans += 1LL * (sz[find(i, 0)][0] - 1) * (sz[find(i, 1)][1] - 1);
    }
    printf("%lld\n", ans);
    return 0;
}
最后修改:2019 年 06 月 01 日 05 : 54 PM

发表评论