[BOJ/Python] 15809_전국시대

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

15809번: 전국시대

문제

전국시대엔 N개의 국가가 존재한다. 각 국가는 1부터 N까지의 번호를 가지고 있다.

또한, 모든 국가는 각자 자신의 국가의 힘을 상징하는 병력을 가지고 있다. 이때 M개의 기록이 주어진다. 각각의 기록은 다음과 같다.

1. 동맹 - 두 나라가 서로 동맹을 맺는다. 두 나라의 병력이 하나로 합쳐진다.
2. 전쟁 - 두 나라가 서로 전쟁을 벌인다. 병력이 더 많은 나라가 승리하며 패배한 나라는 속국이 된다. 이때 남은 병력은 승리한 나라의 병력에서 패배한 나라의 병력을 뺀 수치가 된다. 두 나라의 병력이 같을 경우 두 나라 모두 멸망한다.

모든 나라는 정직하기 때문에 내 동맹의 동맹도 나의 동맹이고, 내 동맹이 적과 전쟁을 시작하면 같이 참전한다. 속국인 경우도 동맹의 경우와 마찬가지이다.

따라서, 전쟁에서 진 국가와 동맹인 다른 국가 또한 전쟁에서 이긴 국가의 속국이 된다.

모든 기록이 끝났을 때 남아있는 국가의 수를 출력하고, 그 국가들의 남은 병력의 수를 오름차순으로 출력하는 프로그램을 작성하시오.

단, 여러 국가가 서로 동맹이거나 속국 관계인 경우는 한 국가로 취급한다.

입력

첫 번째 줄에 국가의 수를 나타내는 N과 기록의 수 M이 주어진다. (1 ≤ N, M ≤ 100,000)

두 번째 줄 부터 N개의 줄에 걸쳐 i번째 국가의 병력 Ai (1 ≤ i ≤ N)가 자연수로 주어진다. (1 ≤ Ai ≤ 10,000)

다음 M개의 줄에는 기록이 3개의 정수 O, P, Q로 주어진다. O가 1인 경우 P, Q가 서로 동맹을 맺었음을 의미하고, O가 2인 경우 P, Q가 서로 전쟁을 벌였음을 의미한다.

동맹끼리 다시 동맹을 맺거나 전쟁하는 입력은 주어지지 않는다.

출력

첫째 줄에 남아있는 국가의 수를 출력한다.

다음 줄에 각 국가의 남은 병력의 수를 띄어쓰기를 간격으로 오름차순으로 출력한다.

예제 입력 1

5 3
10
20
30
40
50
1 1 2
1 3 4
2 1 3

예제 출력 1

2
40 50

---

풀이

 

union&find 문제이다.

 

1. 각 국가마다 병력이 존재한다.
2. 1 a b 의 입력이 들어오면, a와 b의 병력을 합친다.
3. 2 a b 의 입력이 들어오면, a와 b 중 더 병력이 큰 쪽의 경우, 상대 국가의 병력의 수만큼 감소하게 된다.
4. 만약 2 a b 의 입력이 들어왔을 때, a와 b의 병력이 같다면, 두 나라는 모두 멸망한다.

코드는 다음과 같다.

 

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,o, a,b):
    a = find_parent(parent,a)
    b = find_parent(parent,b)
 
    if o==1:
        if a<b:
            power[a] += power[b]
            parent[b] = a
        elif a>b:
            power[b] += power[a]
            parent[a] = b
    else:
        if power[a] > power[b]:
            power[a] = power[a] - power[b]
            parent[b] = a
        elif power[a] < power[b]:
            power[b] = power[b] - power[a]
            parent[a] = b
        else:
            parent[a] = 0
            parent[b] = 0
 
 
n,m = map(int,sys.stdin.readline().split())
parent = [i for i in range(n+1)]
power = [0]
 
for _ in range(n):
    power.append(int(sys.stdin.readline()))
 
for _ in range(m):
    o,p,q = map(int,sys.stdin.readline().split())
    union_parent(parent, o, p, q)
 
country = set()
for i in range(n+1):
    country.add(find_parent(parent, i))
 
result = []
for c in country:
    if c!=0:
        result.append(power[c])
 
result.sort()
 
print(len(country)-1)
print(*result)