[BOJ/Python] 25187_고인물이 싫어요

문제 출처:https://www.acmicpc.net/problem/25187

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)가 주어진다. 

출력

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

예제 입력 1

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

예제 출력 1

0
1
1

첫번째 쿼리는 1번 물탱크와 연결된 물탱크는 2번 물탱크 하나이므로, 고인물이 담긴 물탱크와 청정수가 담긴 물탱크의 수가 같다.

두번째 쿼리는 5번 물탱크와 연결된 물탱크는 3, 4번 물탱크이므로, 고인물이 담긴 물탱크보다 청정수가 담긴 물탱크가 더 많다. 

예제 입력 2

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

예제 출력 2

1
0
1

---

풀이

 

1. N개의 물탱크가 존재한다. (N개의 정점이 존재한다.)
2. 물탱크는 각각 청정수 혹은 고인물이 담겨 있다. (청정수는 1, 고인물은 0이다.)
3. 물탱크들은 M개의 파이프를 통하여 연결되어 있다. (M개의 간선이 존재한다.)
4. 이때, 연결된 물탱크들 중 청정수가 담긴 물탱크의 수가 고인물이 담긴 물탱크의 수보다 많아야 청정수를 얻을 수 있다.
5. 주어진 물탱크들을 방문하였을 때, 청정수를 얻을 수 있는지 출력하는 문제이다.

 

풀이는 다음과 같다.

1. 우선 물탱크들이 M개의 파이프를 통하여 연결되어있음을 구현하기 위하여, union&find를 이용하였다.
2. 또한, union&find를 이용하여 구한 각 집합에 속하는 물탱크들 중 청정수가 더 많음을 나타내어야 했다.
3. 이를 나타내기 위하여, 문제에서 청정수를 1, 고인물을 0으로 주어졌지만 청정수를 1, 고인물을 -1로 바꾸었다.
4. 그리고 각 물탱크들을 연결할 때, 기존의 같은 집합에 속해있지 않다면, 해당 집합들의 청정수의 값과 고인물의 값들을 더해주었다. (각 집합의 루트 노드에 해당 집합의 청정수의 값의 합과 고인물의 값의 합을 저장하였다.)
5. 그 결과로 해당 값이 양수라면, 청정수가 고인물보다 많으므로 청정수를 획득할 수 있고, 0 또는 음수라면 청정수가 고인물보다 작거나 같으므로 청정수를 획득할 수 없다.

코드는 다음과 같다.

 

import sys
sys.setrecursionlimit(1000000)
 
def find_parent(parent,x):
    if parent[x] != x:
        parent[x] = find_parent(parent,parent[x])
    return parent[x]
 
def union_parent(parent,a,b):
    a = find_parent(parent,a)
    b = find_parent(parent,b)
 
    if a<b:
        water_tank[a] += water_tank[b]
        parent[b] = a
    elif a>b:
        water_tank[b] += water_tank[a]
        parent[a] = b
 
 
n,m,q = map(int,sys.stdin.readline().split())
 
parent = [i for i in range(n+1)]
 
water_tank = [0]
water_tank.extend(list(map(int,sys.stdin.readline().split())))
for i in range(1,n+1):
    if water_tank[i] == 0:
        water_tank[i] = -1
 
for _ in range(m):
    u,v = map(int,sys.stdin.readline().split())
    union_parent(parent,u,v)
 
for _ in range(q):
    k = int(sys.stdin.readline())
    if water_tank[find_parent(parent,k)]>0:
        print(1)
    else:
        print(0)