CF1156D 0-1-Tree

CF1156D 0-1-Tree - 洛谷 | 计算机科学教育新生态 (luogu.com.cn)

考虑对每个点 \(u\) 统计：在点 \(u\) 中转的路径有多少条？

其中一条路径在点 \(u\) 中转的意思为：这条路径在到达点 \(u\) 前访问的均为 \(0\) 边，到达点 \(u\) 后访问的边均为 \(1\) 边。特别地，如果一条路径起点为 \(u\) 在 \(u\) 中转，它经过的边权全为 \(1\) ；如果一条路径终点为 \(u\) 在 \(u\) 中转，它经过的边权全为 \(0\)。

我们一开始先对树 \(T\) 预处理：断开所有边权为 \(1\) 的边后形成的连通块称作一类连通块；断开所有边权为 \(0\) 的边后形成的连通块称作二类连通块。

容易证明 \(u\) 所在一类连通块和 \(u\) 所在二类连通块存在且仅存在一个交点，这个交点就是 \(u\)。

那么，一条路径 \((s, t)\) 在点 \(u\) 中转，充要条件是：

• \(s\) 在 \(u\) 所在的一类连通块中；
• \(t\) 在 \(u\) 所在的二类连通块中；
• \(s, t\) 不同时等于 \(u\)。

如此，\(s \rightsquigarrow u \rightsquigarrow t\) 才是一个合法的在 \(u\) 中转的路径。

因此，若 \(u\) 所在一类连通块大小为 \(p\)，所在二类连通块大小为 \(q\)，那么在 \(u\) 中转的路径有 \(p \times q - 1\) 条。\(-1\) 是扣除 \(s = u = t\) 的情况。

时间复杂度 \(\mathcal{O}(n)\)。

/*
 * @Author: crab-in-the-northeast 
 * @Date: 2022-10-18 16:36:45 
 * @Last Modified by: crab-in-the-northeast
 * @Last Modified time: 2022-10-18 17:02:35
 */
#include <bits/stdc++.h>
#define int long long
inline int read() {
    int x = 0;
    bool flag = true;
    char ch = getchar();
    while (!isdigit(ch)) {
        if (ch == '-')
            flag = false;
        ch = getchar();
    }
    while (isdigit(ch)) {
        x = (x << 1) + (x << 3) + ch - '0';
        ch = getchar();
    }
    if(flag)
        return x;
    return ~(x - 1);
}

const int maxn = (int)2e5 + 5;
typedef std :: pair <int, int> pii;
std :: vector <pii> G[maxn];

int con[2][maxn];
int siz[2][maxn];

void dfs(int u, int x, int id) {
    con[x][u] = id;
    ++siz[x][id];
    for (pii e : G[u]) {
        int v = e.first, w = e.second;
        if (w != x || con[x][v])
            continue;
        dfs(v, x, id);
    }
}


signed main() {
    int n = read();
    for (int i = 1; i < n; ++i) {
        int u = read(), v = read(), w = read();
        G[u].emplace_back(v, w);
        G[v].emplace_back(u, w);
    }

    for (int u = 1; u <= n; ++u) {
        if (!con[0][u])
            dfs(u, 0, u);
        if (!con[1][u])
            dfs(u, 1, u);
    }

    int ans = 0;
    for (int u = 1; u <= n; ++u)
        ans += siz[0][con[0][u]] * siz[1][con[1][u]] - 1;
    
    printf("%lld\n", ans);
    return 0;
}