[그래프] 그래프 알고리즘 (4) 실전 문제 - 팀 결성

이코테 p.298 실전 문제 - 팀 결성

입력 예시

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의 범위가 모두 최대 100,000이므로 경로 압축 방식의 서로소 집합 자료구조를 이용하여 시간 복잡도를 개선함
• a와 b 각각에 대한 찾기 연산의 수행 결과가 같은지 비교하여 그 결과를 result에 저장 후 출력

# 같은 팀 여부 확인 연산 정의
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 입력 받기
n, m = map(int, input().split())

# 루트 노드를 저장하기 위한 리스트
parent = [0] * (n + 1)
# 자기 자신의 루트노드는 자기 자신으로 초기화
for i in range(1, n + 1):
  parent[i] = i

result = [] # 결과를 저장할 리스느
for _ in range(m):
  oper, a, b = map(int, input().split())
  if oper == 0: # 팀 합치기 연산인 경우
    union_parent(parent, a, b)
  else: # 같은 팀 여부 확인 연산인 경우
    if find_parent(parent, a) == find_parent(parent, b): # 같은 팀인 경우
      result.append('YES')
    else: # 같은 팀이 아닌 경우
      result.append('NO')

# 결과 출력
for i in result:
  print(i)