[CodingTest] 팀 결성 ( Graph 알고리즘 )

 

◎ 문제

 

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

 

◎ 문제풀이

 

이코테 교재는 연산을 하나씩 입력받고 하나씩 결과를 처리했다. 나는 그러지 않고 그래프 문제이므로 입력받은 연산을 모두 그래프 자료구조에 넣어서 한번에 처리하고 싶었다. 그랬더니 문제가 발생했다.

 

UNION-FIND 알고리즘을 사용하므로 재귀호출이 너무 많이 일어나, RecursionError: maximum recursion depth exceeded in comparison 에러가 발생했다. 그래서 재귀호출 가능깊이를 늘려주었더니 Segmentation fault (core dumped) 에러가 발생했다. 재귀호출이 반복되면서 접근 가능한 메모리 범위를 넘어선 것이다.

 

문제는 재귀호출이므로 재귀호출을 스택자료구조를 활용하여 반복문으로 바꾸어 주었더니 문제가 해결되었다. 

 

 

◎ 코드 ( python )

#재귀호출을 이용함 FIND 메소드 ( 사용X )
def find_root_recursive(parent,x) :
  if parent[x] != x : parent[x] = find_root_recursive(parent,x)
  return parent[x]

#STACK + 반복문을 활용한 FIND 메소드 
def find_root_while(parent,x) :
  stack = []
  while parent[x] != x :
    stack.append(x)
    x = parent[x]

  for node in stack :
    parent[node] = x

  return parent[x]

#UNION 메소드 
def union_graph(parent,a,b) :
  a_root = find_root_while(parent,a)
  b_root = find_root_while(parent,b)

  if a_root > b_root : parent[a_root] = b_root
  else : parent[b_root] = a_root


n,m = map(int,input().split()) # n : 노드의 개수, m : 연산의 개수
parent = [0]*(n+1) # 부모 테이블 생성
graph= [[] for i in range(n+1)] # 그래프 생성

#부모 테이블 초기화
for i in range(1,n+1) :
  parent[i] = i

#그래프 입력받기
for i in range(m) :
  e,a,b = map(int,input().split())
  graph[a].append((e,b)) 

#문제 해결하기
for i in range(1,n+1) :
  for node in graph[i] :
    a = i
    e,b = node
    #합치기
    if e == 0 :
      union_graph(parent,a,b)
    #탐색
    elif e == 1 :
      a_root = find_root_while(parent,a)
      b_root = find_root_while(parent,b)
    
      if a_root == b_root : print("YES")
      else : print("NO")

 

 

---

 

참고자료

이것이 취업을 위한 코딩 테스트다 with 파이썬 : 네이버 도서