[Python] 백준 1717번 집합의 표현

문제

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

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

입력

첫째 줄에 n(1 ≤ n ≤ 1,000,000), m(1 ≤ m ≤ 100,000)이 주어진다. m은 입력으로 주어지는 연산의 개수이다. 다음 m개의 줄에는 각각의 연산이 주어진다. 합집합은 0 a b의 형태로 입력이 주어진다. 이는 a가 포함되어 있는 집합과, b가 포함되어 있는 집합을 합친다는 의미이다. 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산은 1 a b의 형태로 입력이 주어진다. 이는 a와 b가 같은 집합에 포함되어 있는지를 확인하는 연산이다. a와 b는 n 이하의 자연수 또는 0이며 같을 수도 있다.

출력

1로 시작하는 입력에 대해서 한 줄에 하나씩 YES/NO로 결과를 출력한다. (yes/no 를 출력해도 된다)

1717번: 집합의 표현

---

[나의 풀이]

getParents 함수의 경로 최적화가 필요한 문제이다. (참고설명은 아래 링크 걸어둠)

최악의 경우는 getParents가 O(N)이라서 총 (M*N)이 될 수 있으니 부모노드를 최상위계층으로 업데이트하는게 포인트

이런걸 경로 압축(Path Compression) 이라고 하고, 시간 복잡도: O(logN)이 된다고 한다.

import sys
sys.setrecursionlimit(100000000)

def getParents(a):
    # 종료 조건
    if parents[a] == a:
        return a
    parents[a] = getParents(parents[a]) #경로 최적화
    return parents[a]

def union(a, b):
    a = getParents(a)
    b = getParents(b)
    if a == b: return
    if a < b: parents[b] = a
    else: parents[a] = b

def findUnion(a, b):
    if getParents(a) == getParents(b):
        print("YES")
    else:
        print("NO")

n, m = map(int, sys.stdin.readline().split())
parents = list(range(n+1)) # 노드 갯수의 +1로, Node 0~1~7
for _ in range(m):
    q, a, b = map(int, sys.stdin.readline().split())
    if q == 0:
        union(a, b)
    else:
        # 같은 부모인지 확인해서 출력
        if getParents(a) == getParents(b):
            print("YES")
        else:
            print("NO")

https://gmlwjd9405.github.io/2018/08/31/algorithm-union-find.html

[알고리즘] Union-Find 알고리즘 - Heee's Development Blog