[BOJ/Python] 3830번 문제풀이

문제 설명

상근이는 매일 아침 실험실로 출근해서 샘플의 무게를 재는 일을 하고 있다. 상근이는 두 샘플을 고른 뒤, 저울을 이용해서 무게의 차이를 잰다.

교수님의 마음에 들기 위해서 매일 아침부터 무게를 재고 있지만, 가끔 교수님이 실험실에 들어와서 상근이에게 어떤 두 샘플의 무게의 차이를 물어보기도 한다. 이때, 상근이는 지금까지 잰 결과를 바탕으로 대답을 할 수도 있고, 못 할 수도 있다.

상근이는 결과를 출근한 첫 날부터 공책에 적어 두었다. 하지만, 엄청난 양의 무게가 적혀있기 때문에, 교수님의 질문에 재빨리 대답할 수가 없었다. 이런 상근이를 위해서 프로그램을 만들려고 한다.

상근이가 실험실에서 한 일이 순서대로 주어진다. 어떤 두 샘플의 무게의 차이를 구할 수 있는지 없는지를 알아내는 프로그램을 작성하시오.

입력

입력은 여러 개의 테스트 케이스로 이루어져 있다.

각 테스트 케이스의 첫째 줄에는 샘플의 종류의 개수 N (2 ≤ N ≤ 100,000)과 상근이가 실험실에서 한 일의 수 M (1 ≤ M ≤ 100,000)이 주어진다. 샘플은 1번부터 N번까지 번호가 매겨져 있다. 다음 M개 줄에는 상근이가 실험실에서 한 일이 주어진다.

상근이가 무게를 쟀다면, ! a b w와 같은 형식으로 주어진다. 이 뜻은 b가 a보다 w그램 무겁다는 뜻이다. (a ≠ b) w는 1,000,000을 넘지 않는 음이 아닌 정수이다. 모든 측정은 정확하고, 일관성을 유지한다.

교수님의 질문은 ? a b와 같은 형식으로 주어진다. 이 뜻은 b가 a보다 얼마나 무거운지를 출력하라는 뜻이다.

마지막 줄에는 0이 두 개 주어진다.

출력

교수님의 질문 (? a b)이 입력으로 들어올 때 마다, 지금까지 측정한 결과를 바탕으로 a와 b의 무게 차이를 계산할 수 있다면, b가 a보다 얼마나 무거운지를 출력한다. 무게의 차이의 절댓값이 1,000,000을 넘지 않는다. 만약, 측정한 결과를 바탕으로 무게의 차이를 계산할 수 없다면, "UNKNOWN"을 출력한다.

 

---

풀이 과정

이 문제는 union-find 알고리즘 응용 문제로, 기존의 union-find의 경로압축과정에서 약간의 변형이 필요하다.

따라서 아래와 같은 방법으로 풀이할 수 있다.

1. 모든 노드에 대해서, 현재 노드의 값, 부모노드, 부모노드의 값을 저장한다.
2. union의 경우 a,b,w가 주어질 때, (a와 b의 현재의 차이)-w 만큼을 a에 대해서 갱신하고, a의 부모를 b의 부모에 포함시킨다.
3. 경로압축의 경우, 현재 부모노드와 현재 부모노드의 값을 미리 저장하고, 경로압축을 한 뒤 현재 부모노드의 값이 변했다면 그 변화량만큼을 현재 값에도 갱신해준다.
4. 출력은 a,b의 부모가 같으면 차이를 출력, 다르면 unknown을 출력한다.
 
2,3번의 과정에서 부모노드의 값을 저장해두어 변화량을 체크하는것이 핵심 아이디어이다.

 

유니온 파운드 알고리즘 관련 참고 자료 | https://johoonday.tistory.com/148

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

# 노드 x와 y를 합치고 가중치 w를 반영하여 부모 노드 및 값 갱신
def union(x, y, w):
  value[x] -= w
  parent[x] = y
  parentvalue[x] = value[y]

# x의 최상위 부모 노드를 찾아서 반환하고, 부모 노드의 값을 갱신
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] == "?":
      # ? 쿼리일 경우 값을 출력하거나 "UNKNOWN" 출력
      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])