[이코테] 09. 그래프 이론 - 팀 결성

---

📢 '이것이 코딩 테스트다 with 파이썬' 책을 공부하고 복습하기 위해 작성했습니다.

---

문제

학생들에게 0부터 N번까지의 번호를 부여했다. 처음에는 모든 학생이 서로 다른 팀으로 구분되어 총 N+1개의 팀이 존재한다.

1. '팀 합치기' 연산은 두 팀을 합치는 연산이다.
2. '같은 팀 여부 확인' 연산은 특정한 두 학생이 같은 팀에 속하는 지를 확인하는 연산이다.

M개의 연산을 수행할 수 있을 때, '같은 팀 여부 확인' 연산에 대한 연산 결과를 출력하라

 

입력 조건

• 첫째 줄에 N, M이 주어진다. M은 입력으로 주어지는 연산의 개수이다. (1 ≤ N, M ≤ 100,000)
• 다음 M개의 줄에는 각각의 연산이 주어진다.
• '팀 합치기' 연산은 0 a b 형태로 주어진다. 이는 a번 학생이 속한 팀과 b번 학생이 속한 팀을 합친다는 의미이다.
• '같은 팀 여부 확인' 연산은 1 a b 형태로 주어진다. 이는 a번 학생과 b번 학생이 같은 팀에 속해있는 지를 확인하는 연산이다.
• a와 b는 N 이하의 양의 정수이다.

출력 조건

• '같은 팀 여부 확인' 연산에 대하여 한 줄에 하나씩 YES 또는 'NO'로 결과를 출력한다.

 

입력 예시
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

출력 예시
NO
NO
YES

---

나의 풀이

N, M = map(int, input().split())

# 부모 테이블 (초기 = 본인)
team = list(range(N+1))


# x의 루트 노드 찾기
def find_parent(parent, x):
    if parent[x] != x:
        parent[x] = find_parent(parent, parent[x])
    return parent[x]


# 팀 합치기
def union(parent, x, y):
    x = find_parent(parent, x)
    y = find_parent(parent, y)
    
    # 작은 노드를 큰 노드의 부모로
    if x > y:
        parent[x] = y
    else:
        parent[y] = x


for _ in range(M):
    x, a, b = map(int, input().split())

    # 팀 합치기 수행
    if x == 0:
        union(team, a, b)

    # 같은 팀 여부 확인
    elif x == 1:
        # 루트 노드가 같으면, 같은 팀
        if find_parent(team, a) == find_parent(team, b):
            print('YES')
        else:
            print('NO')

---

해설

서로소 집합 알고리즘 문제로, N과 M의 범위가 모두 최대 100,000이다.

따라서 경로 압축 방식의 서로소 집합 자료구조를 이용하여 시간 복잡도를 개선해야 한다.

 

팀 합치기는 union 연산, 같은 팀 여부 확인은 find 연산으로 수행한다.

---