[BOJ/Python] 17250_은하철도

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

17250번: 은하철도

문제

하나의 은하 안에는 여러 행성들이 존재한다. 문명의 기술 발전으로 은하 내의 모든 행성들은 서로 여행할 수 있게 되었다.

드디어 오늘, 80,000 광년 떨어진 다른 은하와 우리 은하를 연결하는 은하 철도가 개통된다.

은하 철도가 개통되면 더 많은 행성을 여행할 수 있다는 사실에 은하 내 모든 행성의 주민들은 들떠있는 분위기이다.

우주철도공사 G-Express는 앞으로의 은하 철도 계획을 발표하였다.

우주는 너무 넓기 때문에, G-Express사는 은하가 연결될 때마다 몇 개의 행성들이 서로 여행할 수 있게 되었는 지를 알려주고자 한다.

G-Express사 기술개발팀의 직원인 당신에게 이 프로그램의 업무 요청이 들어왔다. 각 은하들의 행성 수와 철도 계획이 주어지면 해당 철도를 이용할 수 있는 행성들의 수를 실시간으로 안내하는 프로그램을 만들자.

입력

첫 번째 줄에 은하의 수 N과 철도의 개수 M이 주어진다.

두 번째 줄부터 N개의 줄에 N개의 각 은하 내에 존재하는 행성들의 수가 1번 은하부터 차례대로 주어진다. (행성을 세는 단위는 조(10^12) 단위이다.)

그리고 N+2 번째 줄부터 M개의 줄에 걸쳐 은하와 은하 사이를 잇는 철도가 주어진다. 같은 은하 사이에 여러 개의 철도가 건설될 수 있다.

입력되는 N은 2 ≤ N ≤ 100,000, M은 1 ≤ M ≤ 100,000이고, 각 은하의 행성 수는 100(조)개를 넘지 않으며 아무 행성도 없는 경우는 없다.

출력

철도가 연결될 때마다 해당 철도를 이용할 수 있는 행성들의 수를 한 줄씩 출력한다.

예제 입력 1

5 4
3
9
10
11
15
1 2
2 3
4 5
4 3

예제 출력 1

12
22
26
48

예제 입력 2

5 4
3
1
4
15
9
1 2
3 1
2 3
2 4

예제 출력 2

4
8
8
23

노트

입력 데이터가 큰 관계로, 빠른 입출력을 사용하는 것을 권장합니다.

---

풀이

 

1. N개의 은하(노드)가 존재하고, 각 은하에는 최소 1개 이상의 행성이 존재한다.
2. M개의 철도(간선)를 이용하여 은하와 은하를 연결한다.
3. 이때, 해당 철도를 이용할 수 있는 행성의 개수를 출력하는 문제이다.

 

풀이는 다음과 같다.

1. 철도를 연결한다는 것은 은하와 은하를 하나의 집합으로 합치는 과정이라고 할 수 있다.
2. 따라서, 해당 철도를 이용할 수 있는 행성의 개수는 그 집합에 존재하는 행성의 개수를 뜻한다.
3. 우선 집합의 루트 노드에 해당 집합에 존재하는 행성의 개수를 저장한다.
4. 두 집합을 합칠 때는 A집합의 행성의 개수와 B집합의 행성의 개수를 더하여, 합쳤을 때의 루트 노드에 저장한다.
5. 출력의 경우, 해당 집합의 루트 노드가 갖고 있는 행성의 개수를 출력한다.

 

코드는 다음과 같다.

 

import sys
sys.setrecursionlimit(int(1e6))
 
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:
        planet[a] += planet[b]
        parent[b] = a
        return planet[a]
    elif a>b:
        planet[b] += planet[a]
        parent[a] = b
    
    return planet[b]
 
 
n,m = map(int,sys.stdin.readline().split())
parent = [i for i in range(n+1)]
planet = [0]
for _ in range(n):
    planet.append(int(sys.stdin.readline()))
 
for _ in range(m):
    a,b = map(int,sys.stdin.readline().split())
    print(union_parent(parent,a,b))

 

• planet[i]은 i집합에 속한 행성의 수를 나타낸다.
• 만약 두 집합을 합치게 되면, 둘 중 루트 노드가 되는 planet에 자손이 되는 planet을 더해준다.
• 자손이 된 planet의 값은 더 이상 사용하지 않고, 출력의 경우 루트 노드의 planet을 출력한다.
• 만약 두 행성이 이미 같은 집합에 속한다면, 그냥 루트 노드의 planet을 출력한다.