Union - Find with Python

📌 문제

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

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

입력

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

출력

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

💡유니온파인드

• 동적 연결 요소(disjoint-set) 문제를 해결하기 위한 알고리즘입니다.
• 이를 통해 특정 요소 특정 요소들이 같은 집합에 속하는지 효율적으로 판별하거나 두 집합을 합치는 연산을 수행할 수 있습니다.
• Disjoint-set 자료구조라고도 불리며, 네트워크 연결, 그룹핑 문제, 경로 압축 문제 등에서 자주 사용됩니다.

🧑‍💻 코드 풀이

Find - 부모 노드 찾기

def find_parent(parent, x):
    if parent[x] != x:
        parent[x] = find_parent(parent, parent[x])
    return parent[x] 

• 주어진 요소가 속한 집합의 대표 노드를 찾는 연산을 합니다.

Union - 집합 합치기

def union_parent(parent, x, y):
    a = find_parent(parent, x)
    b = find_parent(parent, y)

    if a < b:
        parent[b] = a
    else:
        parent[a] = b

• 두 개의 집합을 하나로 합치는 연산을 합니다.

Solution 로직

def solution(n, m):
    parent = [i for i in range(n + 1)]

    for i in range(m):
        flag, x, y = map(int, sys.stdin.readline().split())

        if flag == 0:
            union_parent(parent, x, y)
        else:
            if find_parent(parent, x) == find_parent(parent, y):
                print('yes')
            else:
                print('no')

• union-find 를 이용해 부모가 동일하다면 yes를 출력합니다.
• 부모가 같지 않다면 no를 출력합니다.