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

문제

• 학교에서 학생들에게 0번부터 N번까지 번호를 부여
  • 1 <= N, M <= 100,000
• 처음에는 모든 학생이 서로 다른 팀으로 구분되어, N+1 개의 팀이 존재
• 선생님은 '팀 합치기' 연산과 '같은 팀 여부 확인' 연산을 사용할 수 있다.
  
  • '팀 합치기' : 0 a b 형태
  • '같은 팀 여부 확인' : 1 a b 형태
• 선생님이 M개 연산 수행할 때, '같은 팀 여부 확인' 연산에 대한 연산 결과 출력하기
  • YES 혹은 NO 결과를 출력한다.

 

풀이

• 전형적인 유니온 파인드 문제
• 시간 복잡도 : O( M logN) 이므로 가능하다.

 

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

N, M = map(int, input().split())
parent = [i for i in range(N+1)]

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

    if op == 0:
        union_parent(parent, a, b)
    elif op == 1:
        if find_parent(parent, a) == find_parent(parent, b):
            print("YES")
        else:
            print("NO")

 

References

• '이것이 코딩테스트다' ch10 그래프 이론 - 팀 결성