Algorithm (PS)

[백준] 1717 집합의 표현 Python - union-find 알고리즘

minjiwoo 2022. 9. 27. 15:27
728x90

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

 

1717번: 집합의 표현

첫째 줄에 n(1 ≤ n ≤ 1,000,000), m(1 ≤ m ≤ 100,000)이 주어진다. m은 입력으로 주어지는 연산의 개수이다. 다음 m개의 줄에는 각각의 연산이 주어진다. 합집합은 0 a b의 형태로 입력이 주어진다. 이는

www.acmicpc.net

 

집합을 합친다고 해서 처음엔 단순하게 배열과 배열을 합쳤는데, 이렇게 하면 시간초과가 난다 ! 
그래서 찾은 방법이 union-find로 배열과 배열을 합칠 때, 부모 노드 (root node)를 이용해서 부모 노드가 더 작은 쪽으로 가게끔 합쳤다 

그리고 각각의 원소가 하나의 집합에 속하는지 확인할 때는 부모노드가 동일한지만 확인해주면 되니까 더 효율적으로 구할 수 있다 
오랜만에 union-find 본다

import sys
sys.setrecursionlimit(100000)
n, m = map(int, sys.stdin.readline().split())
parents = [0] * (n+1)

for i in range(1, n+1):
    parents[i] = i # 자기 자신을 부모로 설정

def find(a):
    if parents[a] != a: # 자기 자신이 root 노드가 아닌경우 다시 find 연산 수행하기
        parents[a] = find(parents[a])
    return parents[a]

def union(a, b):
    a = find(a)
    b = find(b)
    if a < b:
        parents[b] = a # 루트 노드가 값이 작은 것이 상위라고 가정하기
    else:
        parents[a] = b

for _ in range(m):
    op, a, b = map(int, input().split())
    if op == 0: # 합치기
        union(a, b)
    else: # 동일 집합인지 판단
        if find(a) == find(b):
            print('YES')
        else:
            print('NO')
728x90