[백준 P25187] 고인물이 싫어요 Java

P25187 고인물이 싫어요

25187번: 고인물이 싫어요

문제

재형이는 청정수를 좋아하고 고인물을 싫어한다. 오늘도 청정수를 구하기 위해 물탱크들이 있는 마을에 방문한다.

마을에는 N개의 물탱크가 존재하고, 각 물탱크는 청정수 또는 고인물을 저장하고 있다. 그리고 물탱크는 공급의 편의를 위해 M개의 파이프로 서로 연결되어 있다.

청정수를 얻기 위해 K번 물탱크에 방문했을 때, K번 물탱크와 K번 물탱크에서 0개 이상의 파이프를 거쳐 이동 가능한 물탱크 중, 청정수가 담긴 물탱크의 수가 고인물이 담긴 물탱크의 수보다 더 많은 경우 청정수를 얻을 수 있다.

방문할 예정인 물탱크에 대한 정보가 주어질 때마다, 청정수를 얻을 수 있는지 구해보자.

입력값

첫째 줄에 물탱크의 수 N(1≤N≤100000)과 파이프의 수 M(0≤M≤200000), 그리고 물탱크에 방문할 횟수 Q(1≤Q≤100000)가 주어진다.

둘째 줄에 1번부터 N번 물탱크까지 순서대로 들어있는 물의 종류가 주어진다. 청정수는 1, 고인물은 0으로 주어진다.

다음 3번째부터 M+2번째 줄까지 파이프가 연결하는 두 물탱크의 번호 u,v(1≤u,v≤N,u≠v)가 주어진다. 같은 두 물탱크를 연결하는 파이프가 여러 개 존재할 수 있다.

M+3번째부터 M+Q+2번째 줄까지 방문할 물탱크의 번호 K(1≤K≤N)가 주어진다.

5 3 3
1 0 1 1 0
1 2
3 4
4 5
1
5
4

5 5 3
1 0 1 1 0
1 2
1 3
1 4
2 3
3 4
1
5
4

출력값

Q개의 줄에 각 물탱크에 방문했을 때 청정수를 얻을 수 있다면 1을, 아니면 0을 주어진 정보 순서대로 출력한다.

0
1
1

1
0
1

알고리즘

물탱크의 수와 물탱크 안에 청정수 혹은 고인물 중 어떤 물이 담겨있는지 정보가 주어진다. 이 문제의 키 포인트는 연결되어 있는 물탱크를 모두 방문하여 청정수의 수가 고인물의 수보다 많을 때만 청정수를 얻게 된다는 점이다. 문제에는 청정수를 1, 고인물을 0으로 나타낸다. 청정수와 고인물을 효과적으로 비교하기 위해 고인물의 경우에는 -1으로 바꾸어 저장한다.

그래서 청정수가 2, 고인물인 1개 연결된 물탱크에서는 청정수가 2 + (-1) = 1 개 더 많으므로 청정수를 얻을 수 있다. 하지만 청정수가 2, 고인물이 2개 연결된 물탱크에서는 2 + (-2) = 0이므로 청정수를 얻을 수 없다. 즉 물탱크의 연결 정보가 주어지면, 각 집합의 대표값 (가장 작은 값)의 인덱스에 지금까지 방문한 모든 물탱크의 청정수와 고인물 값을 더하여 저장한다. 그 값이 1 이상이라면 청정수를 얻게되고, 1 미만이라면 청정수를 얻을 수 없다.

parent 는 각 물탱크의 대표값을 저장해둔 배열이다.

whichWater 은 각 물탱크와 연결된 물탱크들을 다 탐색했을 때 얻게되는 청정수 + 고인물의 값 이다.

parent = new int[N + 1];
whichWater = new int[N + 1];

for (int i = 1 ; i <= N ; i++) parent[i] = i;

st = new StringTokenizer(br.readLine());
for (int i = 1; i <= N; i++) whichWater[i] = Integer.parseInt(st.nextToken()) == 1 ? 1 : -1;

find는 전형적인 union-find 알고리즘의 find 처럼 구현하면 된다. 매개변수로 들어오는 숫자가 그 집합의 대표값이라면 반환하고 그렇지 않다면 한번 재귀 호출을 하여 같은 집합의 대표값(가장 작은 값)을 찾는다.

public static int find(int target) {
    if (parent[target] == target) return target;
    else return parent[target] = find(parent[target]);
}

이 문제의 핵심은 union이다. 두 물탱크를 새로 연결시키면, 두 물탱크 집합에 속한 청정수 + 고인물 값(whichWater에 들어있는 값)이 업데이트 된다. 따라서 집합의 대표값(parent)를 업데이트 하면서 whichWater도 업데이트 해야한다. 지금 대표값은 집합 내에서 가장 작은 값을 의미한다. 따라서 a 가 속한 집합의 대표값 aParent 와 b가 속한 집합의 대표값 bParent의 대소를 비교하여 더 큰 값을 가지는 집합의 대표값 (parent) 를 더 작은 집합의 대표값으로 바꿔야 한다.

만약 (1, 3)을 연결한다고 가정하자. 여기서 3번 물탱크에 방문했다고 생각해보자. 3번 물탱크는 1번 물탱크와 연결되어 있다. 그리고 (1, 3)의 대표값은 1이기 때문에 parent[3]의 값은 1이 된다. 여기서 그 집합의 청정수 + 고인물 값을 구하려면 whichWater[1]과 whichWater[3]을 더하면 된다. 지금은 두 개의 물탱크만 연결되어 간단한 연산으로 구할 수 있지만, 여러 개의 물탱크를 연결하면 복잡한 연산이 된다. 따라서 대표값을 업데이트 시켜줄 때 whichWater의 값도 업데이트 시켜주어 연결된 모든 물탱크들의 청정수 + 고인물 값을 구할 수 있도록 할 것이다.

1번과 연결되는 모든 물탱크의 대표값 (parent)는 1일 것이다. 따라서 union 연산을 통해 대표값을 바꿔 줄 때마다 whichWater[1] 에 값을 누적하여 더해주면 최종적으로 연결된 모든 물탱크의 청정수 + 고인물의 값을 구할 수 있다.

결과적으로 문제에 제시된 조건에서 (1, 3)을 연결할 때, parent와 whichWater은 다음의 값으로 변경된다.

1과 3이 연결된 이후에 업데이트 된 값을 보자.

만약 1번 물탱크에 방문한다면 1번 물탱크의 대표값은 1이기 때문에 whichWater[1]을 조회하면 연결된 물탱크들의 청정수 + 고인물 값을 얻는다.

3번 물탱크에 방문할 경우에는 물탱크의 대표값(parent[3])은 1이므로 whichWater[1]을 조회하면 연결된 물탱크들의 청성수 + 고인물 값을 얻는다.

public static void union(int a, int b) {
    int aParent = find(a);
    int bParent = find(b);

    if (aParent < bParent) {
        parent[bParent] = aParent;
        whichWater[aParent] += whichWater[bParent];
    }
    else if (bParent < aParent) {
        parent[aParent] = bParent;
        whichWater[bParent] += whichWater[aParent];
    }
}

구현한 union-find 메서드를 이용한다면 문제를 풀 수 있다. 방문한 물탱크의 whichWater 값이 0 초과면 청정수를 얻을 수 있다고 판단한다. 전체 코드는 다음과 같다.

코드

import java.util.*;
import java.io.*;

public class Main {
    static int N, M, Q;
    static int[] parent;
    static int[] whichWater;

    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());

        N = Integer.parseInt(st.nextToken());
        M = Integer.parseInt(st.nextToken());
        Q = Integer.parseInt(st.nextToken());

        parent = new int[N + 1];
        whichWater = new int[N + 1];

        for (int i = 1 ; i <= N ; i++) parent[i] = i;

        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= N; i++) whichWater[i] = Integer.parseInt(st.nextToken()) == 1 ? 1 : -1;

        for (int i = 0 ; i < M ; i++) {
            st = new StringTokenizer(br.readLine());
            int from = Integer.parseInt(st.nextToken());
            int to = Integer.parseInt(st.nextToken());

            if (find(from) != find(to)) union(from, to);
        }

        for (int i = 0; i < Q; i++) {
            int next = Integer.parseInt(br.readLine());
            System.out.println(whichWater[find(next)] > 0 ? 1 : 0);
        }
    }

    public static void union(int a, int b) {
        int aParent = find(a);
        int bParent = find(b);

        if (aParent < bParent) {
            parent[bParent] = aParent;
            whichWater[aParent] += whichWater[bParent];
        }
        else if (bParent < aParent) {
            parent[aParent] = bParent;
            whichWater[bParent] += whichWater[aParent];
        }
    }

    public static int find(int target) {
        if (parent[target] == target) return target;
        else return parent[target] = find(parent[target]);
    }
}