백준 25187 - 고인물이 싫어요

문제

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

리뷰

유니온 파인드로 풀이할 수 있는 간단한 문제였다.

연결된 물탱크 집합내에서 청정수가 담긴 물탱크 수가 고인물이 담긴 물탱크 수보다
많을 때 청정수를 얻을 수 있다. 이 조건을 표현하기 위해 cost라는 배열을 정의하였다.
cost에는 해당 노드를 루트로 하는 분리 집합내에서 청정수 탱크 개수 - 고인물 탱크 개수
값이 저장된다. 이를 위해 고인물 탱크의 cost 초기값을 -1로 설정해주었으며, 유니온 파인드
연산을 수행하며 유니온 연산이 수행될 두 노드의 cost 값을 더해 루트가 되는 노드의
cost를 갱신해주어 답을 도출하였다.

로직의 시간복잡도는 유니온 파인드 연산이 수행되는 부분의 $O(Ma(N))$으로 수렴하며
이는 $M=200,000$, $N=100,000$인 최악의 경우에도 무난히 제한 조건 1초를 통과한다.

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.StringTokenizer;

import static java.lang.Integer.parseInt;

public class Main {
    static int[] cost;
    static int[] parent;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int N = parseInt(st.nextToken());
        int M = parseInt(st.nextToken());
        int Q = parseInt(st.nextToken());

        cost = new int[N + 1];
        parent = new int[N + 1];
        Arrays.fill(parent, -1);

        st = new StringTokenizer(br.readLine());
        int c;
        for (int i = 1; i < cost.length; i++) {
            c = parseInt(st.nextToken());
            cost[i] = c == 1 ? 1 : -1;
        }

        int u, v, r1, r2;
        while (M-- > 0) {
            st = new StringTokenizer(br.readLine());
            u = parseInt(st.nextToken());
            v = parseInt(st.nextToken());

            r1 = find(u);
            r2 = find(v);

            if (r1 == r2) continue;

            if (parent[r1] < parent[r2]) {
                parent[r1] += parent[r2];
                parent[r2] = r1;

                cost[r1] = cost[r1] + cost[r2];
            } else {
                parent[r2] += parent[r1];
                parent[r1] = r2;

                cost[r2] = cost[r2] + cost[r1];
            }
        }

        StringBuilder sb = new StringBuilder();
        while (Q-- > 0) {
            u = parseInt(br.readLine());
            r1 = find(u);

            if (cost[r1] > 0)
                sb.append(1).append("\n");
            else
                sb.append(0).append("\n");
        }

        System.out.print(sb);
        br.close();
    }

    static int find(int u) {
        if (parent[u] < 0) return u;

        return parent[u] = find(parent[u]);
    }
}

결과