[BOJ/Python] 25587_배수로

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

25587번: 배수로

문제

ChAOS 나라에는 총 N개의 도시가 있고 각각 1, 2, 3, …, N번 도시라고 부른다. ChAOS 나라에 각 도시에는 홍수를 막기 위해 배수로가 설치되어 있다. i번 도시의 배수로는 강수량이 A_i이하일 때만 홍수를 막을 수 있다. 추가로 한 도시에만 폭우가 올 때를 대비해, 두 개의 도시를 정해서 양쪽 도시의 배수로 용량을 공유할 수 있는 공사를 하기로 했다. 예를 들어 1번 도시와 2번 도시에 공사를 하고 난 후, 1번 도시와 2번 도시의 강수량의 합이 A_1 + A_2이하라면 1, 2번 도시 모두에 홍수가 나는 것을 막을 수 있고, 그렇지 않다면 1, 2번 도시 모두에 홍수가 나게 된다. 그 후 2, 3번 도시에도 공사를 하면, 세 도시의 강수량의 합이 A_1 + A_2 + A_3이하라면 1, 2, 3번 도시 모두에 홍수가 나는 것을 막을 수 있고, 그렇지 않다면 1, 2, 3번 도시 모두에 홍수가 나게 된다.

그리고 현재 ChAOS 나라에는 전국적으로 폭우가 오고 있다. 현재 i번 도시의 강수량은 B_i다. 여기서 두 가지의 쿼리를 처리하는 프로그램을 작성하자.

•  1 x y : x번 도시와 y번 도시에 공사를 한다.
•  2 : 현재 상태에서 홍수가 날 도시의 개수를 출력한다.

단, 2번 쿼리는 최소 한 개 주어진다.

입력

첫 번째 줄에 도시의 개수인 정수 N (3 ≤ N ≤ 100,000)과 쿼리의 개수인 정수 M (1≤M≤100,000)이 주어진다.

두 번째 줄에는 i번 도시의 배수로 용량을 의미하는 N개의 정수 A_1, A_2, A_3,..., A_N이 주어진다. (0≤A_i≤1000)

세 번째 줄에는 i번 도시의 강수량을 의미하는 N개의 정수 B_1, B_2, B_3,..., B_N이 주어진다. (0≤B_i≤1000)

네 번째 줄부터 M + 3번째 줄까지는 1 x y 또는 2 형태의 쿼리 M개가 한 줄에 하나씩 주어진다. (1≤x,y≤N)

출력

각각의 2번 쿼리마다 정답을 한 줄에 하나씩 출력한다.

예제 입력 1

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

예제 출력 1

2
1

예제 입력 2

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

예제 출력 2

3
0

---

풀이

 

1. 연결되어 있는 도시들의 배수로 용량이 강수량의 합보다 작으면 그 도시 모두에 홍수가 나게 된다.
2. 2의 입력이 들어올 때, 홍수가 나는 도시의 수를 출력하는 문제이다.

우선 코드는 다음과 같다.

 

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,result):
    a = find_parent(parent,a)
    b = find_parent(parent,b)
 
    if a<b:
        if state[b]<0:
            if state[a]+state[b]>=0:
                result-=country[b]
        else:
            if state[a]+state[b]<0:
                result+=country[b]
        
        if state[a]<0:
            if state[a]+state[b]>=0:
                result-=country[a]
        else:
           if state[a]+state[b]<0:
                result+=country[a]
 
        state[a] += state[b]
        country[a] += country[b]
        parent[b] = a
 
    elif a>b:
        if state[b]<0:
            if state[a]+state[b]>=0:
                result-=country[b]
        else:
            if state[a]+state[b]<0:
                result+=country[b]
        
        if state[a]<0:
            if state[a]+state[b]>=0:
                result-=country[a]
        else:
           if state[a]+state[b]<0:
                result+=country[a]
 
        state[b] += state[a]
        country[b] += country[a]
        parent[a] = b
 
    return result
 
 
n,m = map(int,sys.stdin.readline().split())
parent = [i for i in range(n+1)]
country = [1 for _ in range(n+1)]
 
state = [0]
state.extend(list(map(int,sys.stdin.readline().split())))
 
precipitation = [0]
precipitation.extend(list(map(int,sys.stdin.readline().split())))
 
result = 0
for i in range(1,n+1):
    state[i] = state[i] -precipitation[i]
    if state[i]<0:
        result+=1
 
for _ in range(m):
    q = sys.stdin.readline().split()
    if q[0] == '2':
        print(result)
    else:
        x = int(q[1])
        y = int(q[2])
        result = union_parent(parent,x,y,result)

1. 우선 입력의 경우, 배수로의 용량과 강수량을 빼서 state라는 리스트를 만든다.
2. 만약 연결되어 있는 도시의 루트 노드가 i번 노드일 때, state[i]의 값이 음수라면, i번 도시와 연결되어있는 도시들은 모두 홍수가 난다.
3. result는 홍수가 나는 도시의 수를 출력하기 위한 변수이다. 우선 처음 상태는 연결되어있는 도시가 없으므로 state의 값이 음수인 도시들의 개수를 저장한다.
4. 이후, 2의 입력이 들어오면 result를 출력한다.
5. 1 x y의 입력이 들어오면, union_parent 함수를 이용하여 도시들을 하나로 연결한다.
6. 두 도시 a와 b의 루트 노드인 a`과 b`을 구한다.
7. 여기서 country [a`]은 a`과 연결되어 있는 도시의 개수이다. state [a`]은 그 도시의 배수로의 총 용량 - 강수량의 총량의 값으로, 음수라면 그 도시들은 홍수가 나 있는 상태이다.
8.  a`과 b`을 합칠 때, 홍수가 나던 도시가 합침으로써 홍수가 안 날 수도 있고, 홍수가 안 나던 도시가 합침으로써 홍수가 날 수도 있다.
9. 이 경우 result에서 해당 도시의 개수들을 더하고 빼줌으로써 홍수가 난 도시의 개수를 계산할 수 있다.