[백준 3830번] 교수님은 기다리지 않는다 (Python3)

import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**6)

def union(x,y,w):
  value[x] -= w
  parent[x] = y
  parentvalue[x] = value[y]

def findparent(x):
  if parent[x] != x:
    p,pv = parent[x],parentvalue[x]
    parent[x] = findparent(p)
    value[x] += value[p]-pv
    parentvalue[x] = value[parent[x]]
  return parent[x]

while 1:
  N,M = map(int,input().split())
  if not N:
    break
  parent = [i for i in range(N+1)]
  value = [0]*(N+1); parentvalue = [0]*(N+1)
  
  for _ in range(M):
    Q = input().split()
    if Q[0]=="?":
      a,b = map(int,Q[1:])
      print(value[b]-value[a] if findparent(a)==findparent(b) else "UNKNOWN")
    else:
      a,b,w = map(int,Q[1:])
      if findparent(a) != findparent(b):
        union(parent[a],parent[b],w-value[b]+value[a])

간단한 union-find 응용문제.

기존의 union-find의 경로압축과정에서 약간의 변형이 필요하다.

 

풀이를 설명하면,

1. 모든 노드에 대해서, 현재 노드의 값, 부모노드, 부모노드의 값을 저장한다.

2. union의 경우 a,b,w가 주어질 때, (a와 b의 현재의 차이)-w 만큼을 a에 대해서 갱신하고, a의 부모를 b의 부모에 포함시킨다.

3. 경로압축의 경우, 현재 부모노드와 현재 부모노드의 값을 미리 저장하고, 경로압축을 한 뒤 현재 부모노드의 값이 변했다면 그 변화량만큼을 현재 값에도 갱신해준다.

4. 출력은 a,b의 부모가 같으면 차이를 출력, 다르면 unknown을 출력한다.

 

2,3번의 과정에서 부모노드의 값을 저장해두어 변화량을 체크하는것이 핵심 아이디어이다.

 

 

오늘의 교훈) 파이썬에서 재귀함수를 쓸때 항상 recursion error를 조심하자