[이코테] 그래프 알고리즘 - 팀 결성

해당 포스팅의 문제의 출처는 나동빈님의 이것이 취업을 위한 코딩 테스트 교재를 공부하면서 풀이할 때 본인의 사고 과정 흐름과 문제 풀이를 기록하기 위함 입니다. 

 

---

문제설명

학교에서 학생들에게 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로 결과를 출력한다.

사고과정

• 문제에서 '팀 합치기' 와 '같은 팀 여부 확인'의 키워드로 서로소 집합 자료구조의 Union, Find 연산임을 눈치챘다. 그래서 부모 테이블을 자기 자신 노드로 초기화하고 find, Union 연산 함수 정의 후 입력되는 조건에 맞게 find, union 연산을 처리했다. 만약 find 연산을 수행하라고 했을 때는 서로의 부모노드가 같은지 여부에 따라 같은 팀인지 여부를 판별할 수 있었다.
• 확실히 알고리즘 개념을 알고 있는 것이 문제 해결 아이디어를 떠올리는 데 도움을 받았다.

풀이

책 풀이와 거의 유사하므로 필자의 풀이만 제시하겠다.

 

import sys

n, m = map(int, input().split())
parent = [0] * (n+1)
for i in range(n+1):
    parent[i] = i

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

def union_parent(parent, a, b):
    a = find_parent(parent, a)
    b = find_parent(parent, b)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b

for _ in range(m):
    calc, a, b = map(int, input().split())
    if calc == 0:  # 합치기
        union_parent(parent, a, b)
    else:  # 찾기
        root_a = find_parent(parent, a)
        root_b = find_parent(parent, b)
        if root_a == root_b: # 부모노드 같으면 같은팀
            print('YES')
        else:
            print('NO')