백준1717 - 집합의 표현 (Find-Union) Python

https://www.acmicpc.net/problem/1717

해당 문제는 집합의 갯수 N과 연산의 갯수 M, 그리고 연산이 주어진다. 이때, 맨 앞 정수가 0인 경우 두집합을 합치고, 1인 경우는 같은 집합에 있는지 확인한다고 할 때, 1로 시작하는 입력에 대해 원소 a,b가 같은 집합에 포함되어 있으면 'YES' 그렇지 않다면 'NO'를 반환하는 문제다.

 

문제 해결 방식

집합이 각자 자기 자신을 가지고 시작하고, 집합을 합치거나 공통 루트를 찾는다는 점에서, 이전에 크루스칼 알고리즘에서 학습했던 find-union 알고리즘을 떠올려 적용시켰다. 두 집합을 합치는 연산에는 union을, 같은 집합에 포함되어 있는지 확인하는 연산에서는 find를 활용하였다.

 

전체 소스 코드

import sys
data_temp = sys.stdin if sys.platform == 'linux' else open('입력.txt', 'r')
input_data = data_temp.read().splitlines()
sys.setrecursionlimit(150000)

def solution (data) :
  arr = [list(map(int, v.split())) for v in data]
  N,M = arr.pop(0)
  
  roots = {}
  for i in range(N+1) :
    roots[i+1] = i+1
    
  def find (n) :
    if n != roots[n] :
      roots[n] = find(roots[n])
    return roots[n]
  
  def union (x,y) :
    x,y = find(x),find(y)
    if x < y : roots[y] = x
    else : roots[x] = y
  
  result = []
  for order,a,b in arr :
    a,b = a+1, b+1
    if order == 0 : 
      union(a,b)
    else : 
      result.append('NO') if find(a) != find(b) else result.append('YES')
  
  print('\n'.join(result))

solution(input_data)