[그래프 이론] 이코테 (파이썬) 팀 결성 풀이

[문제]

학교에서 학생들에게 0번부터 N번까지의 번호를 부여했다. 처음에는 모든 학생이 서로 다른 팀으로 구분되어,

 

총 N + 1개의 팀이 존재한다. 이때 선생님은 '팀 합치기' 연산과 '같은 팀 여부 확인' 연산을 사용할 수 있다.

 

1. '팀 합치기' 연산은 두 팀을 합치는 연산이다.

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

 

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

 

[입력 조건]

1. 첫째 줄에 N, M이 주어진다. M은 입력으로 주어지는 연산의 개수이다. (1 <= N, M <= 100,000)

 

2. 다음 M개의 줄에는 각각의 연산이 주어진다.

 

3. '팀 합치기' 연산은 0 a b 형태로 주어진다. 이는 a번 학생이 속한 팀과 b번 학생이 속한 팀을 합친다는 의미이다.

 

4. '같은 팀 여부 확인' 연산은 1 a b 형태로 주어진다. 이는 a번 학생과 b번 학생이 같은 팀에 속해 있는지를 확인하는 연산이다.

 

5. 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

 

[풀이]

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 = [0] * (n + 1)

for i in range(0, n + 1) :
  parent[i] = i

for i in range(m) :
  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')

 

이 문제는 서로소 집합 알고리즘을 통해 해결할 수 있다.

 

먼저 n, m을 입력받은 후 부모 테이블(parent)을 초기화 한다. 그 후 부모 테이블상에서 부모를 자기 자신으로 초기화해주고

 

반복문을 통해 각 연산을 하나씩 확인해가며 결과를 가져오는데,

 

연산(oper)이 0일 경우 팀을 합쳐야하기 때문에 union_parent( ) 메서드를 수행한다.

 

union_parent( ) 메서드에서는 특정 원소가 속한 집합을 찾아 두 원소가 속한 집합을 합치는 작업을 진행한다.

 

연산(oper)이 1일 경우 같은 팀에 속해 있는지를 확인해야하기 때문에 find_parent( ) 메서드를 수행한 각 결과값을 비교하여

 

같은 팀에 속해 있다면 'YES'를, 그렇지 않다면 'NO'를 출력하도록 한다.

 

find_parent( ) 메서드에서는 특정 원소가 속한 집합을 찾는데, 만약 해당 노드가 루트 노드가 아니라면 루트 노드를

 

찾을 때까지 재귀적으로 호출하여 수행한다.

 

---

출처

이것이 코딩 테스트다 with 파이썬 - 나동빈 저