백준(baekjoon) 1717번: 집합의 표현(골드 4, 파이썬 PYTHON) - 유니온 파인드

초기에 n+1개의 집합 {0},{1},{2},…,{n}이 있다. 여기에 합집합 연산과, 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산을 수행하려고 한다.

집합을 표현하는 프로그램을 작성하시오.

 

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

1717번: 집합의 표현

 

import sys
sys.setrecursionlimit(1000000)

n,m=map(int,sys.stdin.readline().split())
parent=[None for _ in range(n+1)]

for x in range(n+1):
  parent[x]=x

def find(num):
  if(parent[num]==num):
    return num
  n=find(parent[num])
  parent[num]=n
  return parent[num]

def union(a,b):
  parent_a=find(a)
  parent_b=find(b)
  if(parent_a==parent_b):
    return
  if(parent_a<parent_b):
    parent[parent_b]=parent_a
  else:
    parent[parent_a]=parent_b
 
for _ in range(m):
  num,a,b=map(int,sys.stdin.readline().split())

  if(num==0):
    union(a,b)
  elif(num==1):
    if(find(a)==find(b)):
      print("YES")
    else:
      print("NO")

 

※ 유니온 파인드

 

 

● 만약 두 원소가 연결되어 있으면 더 큰 원소의 부모를 더 작은 수의 원소로 설정해 준다는 것이 기본 전제이다. 우선 모든 원소의 부모를 자기 자신으로 설정해 놓는다.

 

● 합집합 연산이 들어오면 두 원소의 부모를 찾는다. 만약 두 원소의 부모가 같다면 이미 합해져 있는 것이므로 그대로 return 해준다. 만약 한 쪽 원소의 부모가 더 크다면 더 큰 원소의 부모를 더 작은 원소로 설정해준다. 그러면 두 원소의 부모는 같게 되고, 이는 두 원소가 같은 집합에 속해있음을 나타낸다.

 

● 두 원소가 같은 집합에 있는지 확인하는 연산이 들어오면 두 원소의 부모를 찾아준다. 이 때, 원소의 부모가 이 원소가 속해있는 집합에서 가장 작은 원소를 가리키고 있지 않을 수도 있으므로 부모가 자기 자신인 원소가 나올 때 까지 재귀적으로 함수를 호출해 주어야 한다. 최종적으로 찾은 두 원소의 부모가 같으면 YES를 출력해주고 다르면 NO를 출력해준다.