13418번: 학교 탐방하기

https://www.acmicpc.net/problem/13418


$O(n+m)$

가중치가 0인 간선부터 선택해서 mst를 만들 때와 1인 간선부터 선택해서 만들 때를 비교하여 0의 개수 제곱차를 구한다.

#include<cstdio>
int n, m, p[2][1001], r1, r2;
int f(int *a, int x) { return x^a[x] ? a[x] = f(a, a[x]) : x; }
int main() {
    scanf("%d%d", &n, &m);
    for (int i = 0; i <= n; i++) p[0][i] = p[1][i] = i;
    for (int i = 0, x, y, z; i <= m; i++) {
        scanf("%d%d%d", &x, &y, &z);
        p[z][f(p[z], x)] = f(p[z], y);
    }
    r1 = n + 1;
    r2 = -1;
    for (int i = 0; i <= n; i++) r1 -= p[0][i] == i, r2 += p[1][i] == i;
    printf("%d", r1*r1 - r2*r2);
    return 0;
}