[카카오] 시험장 나누기

문제

카카오 인턴을 선발하는 코딩 테스트 시험장이 하나의 이진 트리 형태로 연결되어 있습니다. 아래 그림은 12개의 시험장이 연결된 예시입니다.

1. 하나의 노드는 하나의 시험장을 나타냅니다.
2. 검은 바탕의 흰 숫자는 해당 시험장의 고유 번호(ID)를 나타냅니다.
2-1. 시험장이 n개 있다면, 시험장의 고유 번호는 0부터 n-1까지 부여됩니다.
3. 노드 안의 빨간 숫자는, 해당 시험장의 응시자 수를 나타냅니다.
3-1. 위의 그림에서, 9번 시험장에는 10명, 4번 시험장에는 8명, 6번 시험장에는 20명의 응시자가 시험을 볼 예정입니다.
4. 노드 사이의 간선은 해당 시험장이 연결되어 있음을 의미합니다.
4-1. 위의 그림에서, 9번 시험장은 7번 시험장과, 7번 시험장은 6번 시험장과 연결되어 있습니다.

코딩 테스트를 총괄하는 무지는 안정적인 시험을 위해, 시험장에서 오는 트래픽을 k개의 그룹으로 나누어 각 그룹별 서버로 분산시키기로 하였습니다. 시험장 사이를 연결한 간선들 중 k-1개를 끊어서 시험장을 k 개의 그룹으로 나눌 계획입니다. 이때, 그룹별 최대 트래픽을 최소화하기 위하여 가장 큰 그룹의 인원을 최소화시켜야 합니다.

위의 그림에서 7번과 6번 시험장을 잇는 간선을 끊고, 9번과 7번 시험장을 잇는 간선을 끊는다면, 전체 시험장은 3개의 그룹으로 나누어집니다.

• 주황색 노드로 표시된 A그룹의 인원은 35명(10+8+5+6+1+1+4)
• 보라색 노드로 표시된 B그룹의 인원은 37명(7+30)
• 녹색 노드로 표시된 C그룹의 인원은 40명(20+8+12)

즉, 인원이 가장 많은 그룹은 40명입니다. 다른 어떤 방법으로 시험장을 3개의 그룹으로 나눈다고 해도, 인원이 가장 많은 그룹의 인원이 40명 미만이 되도록 나눌 수는 없습니다.

나눌 그룹의 수를 나타내는 정수 k, 각 시험장의 응시자 수를 나타내는 1차원 정수 배열 num, 시험장의 연결 상태를 나타내는 2차원 정수 배열 links가 매개변수로 주어집니다. 인원이 가장 많은 그룹의 인원이 최소화되도록 k개의 그룹으로 나누었을 때, 최소화된 최대 그룹의 인원을 return 하도록 solution 함수를 완성해주세요.

예시

	k	num	links	result
case 1	3	[12, 30, 1, 8, 8, 6, 20, 7, 5, 10, 4, 1]	[[-1, -1], [-1, -1], [-1, -1], [-1, -1], [8, 5], [2, 10], [3, 0], [6, 1], [11, -1], [7, 4], [-1, -1], [-1, -1]]	40
case 2	1	[6, 9, 7, 5]	[[-1, -1], [-1, -1], [-1, 0], [2, 1]]	27
case 3	2	[6, 9, 7, 5]	[[-1, -1], [-1, -1], [-1, 0], [2, 1]]	14
cas3 4	4	[6, 9, 7, 5]	[[-1, -1], [-1, -1], [-1, 0], [2, 1]]	9

case 1.

문제 본문

case 2.

case 3.

case 4.

조건

• 1 ≤ k ≤ 10,000
• k ≤ num의 길이 ≤ 10,000
  • num\[i\]에는 i번 시험장의 응시자 수가 담겨있습니다.
  • 1 ≤ num의 원소 ≤ 10,000
• links의 길이 = num의 길이
  • links의 i번째 행은 i번 노드(시험장)의 [왼쪽 자식 노드 번호, 오른쪽 자식 노드 번호]입니다.
    • 해당 위치에 자식 노드가 없는 경우 -1이 담겨있습니다.
    • 잘못된 노드 번호나, 하나의 이진 트리 구조가 아닌 입력은 주어지지 않습니다.

답

핵심 아이디어는 최대 그룹 인원 m을 미리 정한 뒤, 이를 만족하는 최소 분할 수 P가 문제에서 제시한 k와 맞아 떨어지는지 확인하는 것입니다.

 

이를 위해서는 크게 두 단계를 구현해야 합니다.

1. m이 주어졌을 때, 그룹 인원이 m을 넘기지 않도록 트리를 분할하고 분할 수를 계산하기
2. m의 범위를 정한 뒤, 이진 탐색을 통해 분할 수 k를 만족하는 m 찾기

0. 루트 찾기

그 전에 입력값에서 트리의 루트를 명시하지 않았으므로, 어느 것에도 참조되지 않는 노드를 찾아 루트로 지정해 주어야 합니다.

    is_root = [1 for _ in num]
    
    for i in range(len(num)):
        l, r = links[i]
        if l >= 0:
            is_root[l] = 0
        if r >= 0:
            is_root[r] = 0
    root = is_root.index(1)

1. 트리 분할

주어진 m에 대하여 트리를 분할하는 작업은 함수 search(m)에서 구현합니다.

 

우선 트리를 순회할 순서를 지정합니다. 순회는 후위순서(postorder, left -> right -> mid)로 진행합니다.

 

order 스택에 후위순서의 반대로(mid -> right -> left) 값을 집어 넣고 마지막에 반전하면 후위순서를 얻을 수 있습니다.

 

※ 트리 재귀를 사용하지 않는 이유: 재귀 함수 호출 시 재귀 스택이 쌓이게 될 텐데, 이진 트리가 항상 균형잡힌 경우는 아니므로 트리가 불균형일 때 스택 한도를 초과할 수도 있습니다.

    def search(m):
        order = []
        stack = [root]
        
        while stack:
            node = stack.pop()
            order.append(node)
            l, r = links[node]
            if l >= 0:
                stack.append(l)
            if r >= 0:
                stack.append(r)
        order.reverse()

재귀 대신 DP를 이용합니다. 각 노드에 대해 [분할 수, 그룹 인원]을 dp 리스트에 저장합니다.

 

후위순서대로 순회를 진행하는데, 각 노드마다 왼쪽자식과 오른쪽 자식에 대해 dp의 값을 참조합니다. 존재하지 않으면, 기본 경우 [0, 0]를 사용합니다. 

    def search(m):
        # ...
        dp = [[0, 0] for _ in num]
        
        for node in order:
            l, r = links[node]
            part_l, cnt_l = dp[l] if l >=0 else [0, 0]
            part_r, cnt_r = dp[r] if r >=0 else [0, 0]
            cnt_c = num[node]
        
            if cnt_c > m or INF in [cnt_l, cnt_r]:
                dp[node] = [INF, INF]
                continue

만약 현재 노드의 인원 수 자체가 m보다 크다면, 현재 경우는 어떻게 해도 분할이 불가능한 경우입니다. 따라서 불가능한 경우라는 의미에서 INF값을 넣어줍니다.

 

만약 왼쪽/오른쪽 자식 중 하나가 INF라면, 똑같이 불가능한 경우이므로 현재 노드도 INF가 됩니다.

    INF = 10 ** 11
    
    def search(m):
        # ...
        dp = [[0, 0] for _ in num]
        
        for node in order:
            # ...
            if cnt_c > m or INF in [cnt_l, cnt_r]:
                dp[node] = [INF, INF]
                continue

 

현재 노드와 자신 사이에서 분할을 결정할 때는 크게 네 가지 경우가 있습니다.

		조건식	part 갱신값	cnt 갱신값
경우 1.	자식과 현재 노드의 구성원을 모두 더해도 m을 넘지 않는 경우	cnt < m	part	cnt
경우 2.	왼쪽 자식과 현재 노드의 구성원을 더했을 때 m을 넘지 않는 경우	cnt_l + cnt_c < m	part + 1	cnt_l + cnt_c
경우 3.	오른쪽 자식과 현재 노드의 구성원을 더했을 때 m을 넘지 않는 경우	cnt_r + cnt_c < m	part + 1	cnt_r + cnt_c
경우 4.	어느 경우에도 m을 넘는 경우	else	part + (l>=0) + (i>=0)	cnt_c

경우 2와 경우 3은 동시에 발생할 수 있습니다. 이 때는 둘 중 한 경우만을 선택해야 합니다.

 

우선 part 갱신값이 더 적은 경우를 선택하며, part 갱신값이 똑같다면 cnt 갱신값이 더 적은 경우를 선택합니다.

    INF = 10 ** 11
    
    def search(m):
        # ...
        dp = [[0, 0] for _ in num]
        
        for node in order:
            # ...
            part = part_l + part_r
            cnt = cnt_l + cnt_r + cnt_c
        
            if cnt <= m:
                dp[node] = [part, cnt]
                continue
        
            cand = []
            if cnt - cnt_l <= m:
                cand.append([part+1, cnt - cnt_l])
            if cnt - cnt_r <= m:
                cand.append([part+1, cnt - cnt_r])
            cand.append([part + int(l>=0) + int(r>=0), cnt_c])
        
            dp[node] = min(cand)
        return dp[root][0]

 

2. 이진탐색

m이 가질 수 있는 최솟값은 각 노드 중에서 최대 인원수 입니다. 노드를 하나 단위로 분할한 경우를 의미합니다.

 

m이 가질 수 있는 최댓값은 모든 노드의 인원수 합입니다. 어떤 연결도 분할하지 않은 경우를 의미합니다.

 

이제 최솟값과 최댓값 사이에서 search(m)의 반환값(=분할 수)을 기준으로 이진탐색을 실시합니다.

    l = max(num)
    r = sum(num)
    answer = r
    while l <= r:
        m = (l + r)//2
        p = search(m)
        if p == INF or p+1 > k:
            l = m + 1
        else:
            r = m - 1
            answer = m
    
    return answer