[BOJ/파이썬] 골드5 1717 - 집합의 표현

 💬 문제

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

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

 

 

🔨 입력

 

 

 

🔨출력 

1로 시작하는 입력에 대해서 a와 b가 같은 집합에 포함되어 있으면 "YES" 또는 "yes"를, 그렇지 않다면 "NO" 또는 "no"를 한 줄에 하나씩 출력한다.

👊🏻 예시 입력 1 

7 8
0 1 3
1 1 7
0 7 6
1 7 1
0 3 7
0 4 2
0 1 1
1 1 1

👊🏻 예시 출력 1

 

NO
NO
YES

 

💡 해결

 

코테에서 유니온 파인드 나왔었는데, 알고리즘은 아는데 구현을 못한게 억울해서 유니온파인드 공부해보자 하고 알고리즘 찾아서 푼 문제

 

find, union에 대한 함수를 각각 생성해주었고, 문제에 따라 쉽게 풀었습니다.

# 특정 원소가 속한 집합을 찾기
def find(x):
    if parent[x] != x:
    	return find(parent[x])
    return parent[x]

다만 find를 구현할때 처음에는 위와같이 구현해주었다. 그런데, 이렇게 되면 parent의 배열들이 엉망진창이 될거고, 그렇게 되면 나중에 입력값이 많아지면 무슨 집합에 속하는지 알기 위해서는 find횟수가 어마어마 할 것이다.

 

이걸 풀면서 경로압축이란것도 처음 알았고, 그를 이용한 find를 이용해서 문제를 풀었다.

👨🏻‍💻  CODE
 

import sys
sys.setrecursionlimit(10**5)
N, M = map(int, sys.stdin.readline().split())
parent = [i for i in range(N+1)]

# 특정 원소가 속한 집합을 찾기
def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])
    return parent[x]

# 두 원소가 속한 집합을 합치기
def union(a, b):
    a = find(a)
    b = find(b)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b


for _ in range(M):
    x, a, b = map(int, sys.stdin.readline().split())
    if x == 0: union(a,b)
    else:
        if find(a) == find(b): print("YES")
        else: print("NO")

 

 

💪🏻 결과

 

경로압축 캐리