[파이썬] 백준 15809 : 전국시대 (골드4)

[파이썬] 백준 15809 : 전국시대 (골드4)

• https://www.acmicpc.net/problem/15809

---

풀이

방향성 생각

• 유니온 파인드 + 조건문

---

전체코드

import sys
input = lambda : sys.stdin.readline().rstrip()

N,M = map(int,input().split())

# 입력 받기
arr = [0]
for _ in range(N):
    arr.append(int(input()))

# 부모 관리
P = [i for i in range(N+1)]
def find(x):
    if P[x] != x:
        P[x] = find(P[x])
    return P[x]

# 입력 쿼리
for _ in range(M):
    cmd,a,b = map(int,input().split())
    pa,pb = find(a),find(b)

    # 부모가 다른 경우에 연산 진행
    if pa != pb:

        # 동맹 : 루트가 작은 쪽으로 설정
        if cmd == 1:
            P[max(pa,pb)] = min(pa,pb)
            arr[min(pa,pb)] += arr[max(pa,pb)]
            arr[max(pa,pb)] = 0

        # 전쟁 : 전쟁에서 승리한 쪽으로 설정
        elif cmd == 2:
            ka,kb = arr[pa],arr[pb]
            if ka>=kb:
                P[pb] = pa
                arr[pa],arr[pb] = ka-kb,0
            else:
                P[pa] = pb
                arr[pa],arr[pb] = 0,kb-ka

answer = [a for a in arr if a>0]
print(len(answer))
print(*sorted(answer))

 

---

코멘트

• 입력 인덱스 1씩 빼서 0~N-1로 관리하면 좋을듯?
• 전쟁 유무도 T/F형태로 받을 수 있어서 좀 더 직관적으로 짤 수 있을듯
• 부모 + 인구관리를 따로 해줘야해서, 기존 유니온파인드 쉬운 문제보다는 실수할거리가 있는 문제.