[#알고리즘/Union-Find/[이코테]: 팀 결성(python)]

안지은·2022년 7월 27일
0
post-custom-banner

Question

학교에서 학생들에게 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

Solution

서로소 집합 알고리즘 문제로, 앞서 배운 서로소 집합 소스코드를 활용하여 문제를 해결한다.

Code

import sys

input = sys.stdin.readline

# 특정 원소가 속한 집합 찾기 (즉, 재귀적으로 부모를 찾으며 원소의 루트 노드(속한 집합) 찾기)
def find_parent(parent, x) :
    if parent[x] != x :
        parent[x] = find_parent(parent, parent[x])
    return parent[x] 
        
# union 연산
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 = [0] * (v + 1)

# 부모 테이블에서, 부모를 자기 자신 번호로 초기화
for i in range(1, n+1) :
    parent[i] = i

# union 연산 각각 수행
for i in range(e) :
    oper, a, b = map(int, input().split())
    if oper == 0 :
        union_parent(parent, a, b)
    elif oper == 1 :
        if find_parent(parent, a) == find_parent(parent, b) :
            print("YES")
        else :
            print("NO")
profile
공부 기록용
post-custom-banner

0개의 댓글