[이것이 코딩 테스트다] 팀 결성 with Python

🔒 문제

 학교에서 학생들에게 0번부터 N번까지의 번호를 부여했다. 처음에는 모든 학생이 서로 다른 팀으로 구분되어, 총 N+1개의 팀이 존재한다. 이때 선생님은 '팀 합치기' 연산과 '같은 팀 여부 확인' 연산을 사용할 수 있다.

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

 선생님이 M개의 연산을 수행할 수 있을 때, '같은 팀 여부 확인' 연산에 대한 연산 결과를 출력하는 프로그램을 작성하시오.

입력

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

✅ 해설

• 전형적인 서로소 집합 찾기 알고리즘 응용 문제
• 같은 팀 여부를 확인할 땐 최상위 부모 노드를 비교
  • 최상위 부모 노드가 같다면 같은 팀
  • 최상위 부모 노드가 다르다면 다른 팀
• 자세한 내용은 https://evolving-developer-new.tistory.com/58 참고

💻 코드

# 학생 수와 연산의 개수 입력받기
n, m = map(int, input().split())

# 부모 노드를 저장하는 배열
parent = []
# 자기 자신으로 부모노드 초기화
for i in range(n+1):
    parent.append(i)

# 최상위 부모 노드를 찾는 함수
def find_top(n):
    if parent[n] != n:
        parent[n] = find_top(parent[n])
    return parent[n]

# 팀 합치기 함수
def merge(n1, n2):
    n1_parent = find_top(n1)
    n2_parent = find_top(n2)
    if n1_parent > n2_parent:
        parent[n1_parent] = n2_parent
    else:
        parent[n2_parent] = n1_parent

# 결과를 저장할 배열
result = []

# 명령 입력받기
for _ in range(m):
    order, n1, n2 = map(int, input().split())
    # 같은 팀 여부 확인
    if order == 1:
        # 최상위 부모 노드 비교
        if find_top(n1) == find_top(n2):
            result.append("YES")
        else:
            result.append("NO")
    # 팀 합치기
    else:
        merge(n1, n2)

for i in result:
    print(i, end=' ')