[Python] 1717 - 집합의 표현 (골드4)

1. 문제 설명

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

1717번: 집합의 표현

 

---

2. 아이디어

내가 떠올린 것은 아니고 union-find 알고리즘을 공부하기 위해 찾아본 문제였다.

애초에 1,000,000개의 정점을 BFS, DFS 같은 완전 탐색으로 구현해서 문제를 풀기엔 무리가 있어 보인다.

 

굉장히 재밌는 알고리즘인데 단순히 리스트에 정점과 간선을 넣는 방식이 아니라

인덱스 번호 자체가 정점을 의미하고 인덱스가 가지는 값은 연결 정보를 의미한다.

 

원래 solution에서 알고리즘 설명은 안 하려 했지만 간단하고 재밌는 알고리즘이라 ㅎㅎ..

인덱스 번호를 정점이라고 치고, list[node]를 조회했더니 어라? 다른 정점의 값을 가리키고 있다면?

바로 둘은 합집합 관계라고 판단할 수 있게 되는 것이다.

 

사실 이것만 설명해도 문제를 푸는 데 지장이 없을 것이다.

 

---

3. 코드

import sys

input = sys.stdin.readline
sys.setrecursionlimit(10**8)

def search_node(idx, _set):
    if idx == _set[idx]:
        return idx
    return search_node(_set[idx], _set)

def union(a, b, _set):
    node1 = search_node(a, _set)
    node2 = search_node(b, _set)

    if node1 == node2:
        return
    if node1 > node2:
        _set[node2] = node1
    else:
        _set[node1] = node2

def solution():
    n, m = map(int, input().split())
    _set = [i for i in range(n+1)]

    for _ in range(m):
        flag, a, b = map(int, input().split())
        if flag == 0:
            union(a, b, _set)
        else:
            if search_node(a, _set) == search_node(b, _set):
                print("YES")
            else:
                print("NO")

solution()