시험장 나누기

• k개의 그룹으로 나눌때의 그룹의 최댓값을 최소화시키는 문제는 파라메트릭 서치의 기본 유형중 하나이다.
• 파라메트릭 서치로 풀 때는, 결정 문제가 그룹의 최댓값을 X로 만들때 필요한 그룹의 갯수가 K 이하인지를 확인하는 것이 된다. 즉, 그룹의 최댓값을 X로 만들수 있는 그룹의 최소 갯수를 구하는 것이 필요하다. 1차원 배열을 이런식으로 그룹들로 나누는 문제들은 흔히 나오는 편이고 이 경우는 최댓값이 고정되있을 때 그룹의 최소 갯수를 그리디한 방법으로 간단하게 구할수 있다. 이 문제에서는 이진트리라서 좀더 복잡하지만 그래도 그리디하게 구할수 있다.
• 어떤 노드 P가 자식노드 C1, C2를 갖는다고 할때, 자식노드들과의 다른 그룹으로 나눠야 할지 말지를 생각해보자. 같은 그룹으로 만들었을때 그룹의 값이 X보다 커질때만 다른 그룹으로 나눈다고 생각하면 된다.
  • 노드 U가 갖고 있는 값을 VAL[U]라고 하고 U를 루트로 하는 서브트리에서 U와 같은 그룹에 속한 노드들의 값의 합을 DP[U]라고 하자.
  • VAL[P] + DP[C1] + DP[C2] ≤ X 라면, 모두 같은 그룹으로 만들면 되고 저 값이 DP[P]가 된다
  • 그렇지 않고, VAL[P] + DP[C1] ≤ X 이고 VAL[P] + DP[C2] > X 라면 C2와 P 사이의 에지를 끊는다. 반대의 경우면 C1과 P사이의 엣지를 끊는다. VAL[P] + DP[C1] ≤ X 이고 VAL[P] + DP[C2] ≤ X 라면.. DP[C1]과 DP[C2]중 작은 값을 가진 차일드를 같은 그룹으로 묶고, 나머지를 다른 그룹으로 나눈다.
  • VAL[P] + DP[C2] > X 이고 VAL[P] + DP[C1] > X 면 별수 없이 양쪽 엣지를 모두 끊어야 한다.
  • 만약 VAL[P] > X 라면 그냥 불가능하다. 만약 이분탐색 범위를 max(VAL[i])부터 시작했다면 이런 경우는 발생하지 않으니까 굳이 조건을 추가하지 않아도 되긴 하다.
• 이렇게 하면 이진 트리에서도 최댓값이 고정되있을 때 그룹의 최소 갯수를 구할수 있다. 자식부터 DP를 값을 채워나가면 된다. 시간복잡도는 O(n)
  • DFS를 하면서 채우면 되긴 하는데.. 문제에서 루트 노드를 안가르쳐주기 때문에 루트노드를 먼저 구해야 한다. 자신을 차일드로 갖는 노드가 하나도 없는 노드를 찾으면 되긴 한다. 다만 그렇게 루트를 찾고 또 매번 결정문제를 풀때마다 DFS를 돌리는게 번거로와서, 그냥 위상정렬과 비슷한 방법으로 리프노드들부터 시작해서 루트 노드에서 끝나는 오더링을 만들고, 그 순서대로 DP값을 계산하는 식으로 처리했다.
• 결정 문제를 해결하는 시간이 O(n)이고, 이분탐색의 범위가 max(VAL[i]) ~ sum(VAL[i]) 로 O(nm)이다. (n: 노드의 수, m: 각노드가 갖을수 있는 최댓값). 따라서 총 시간 복잡도는 O(nlog(nm))

"""Solution code for "Programmers 81305. 시험장 나누기".

- Problem link: https://programmers.co.kr/learn/courses/30/lessons/81305
- Solution link: http://www.teferi.net/ps/problems/programmers/81305
"""

from teflib import algorithm


def solution(k, num, links):

    def can_make_group(group_size):
        dp = [0] * len(num)
        group_count = 1
        for node in nodes:
            c1, c2 = links[node]
            node_val = num[node]
            child_sum1 = dp[c1] if c1 != -1 else 0
            child_sum2 = dp[c2] if c2 != -1 else 0
            if child_sum1 + child_sum2 + node_val <= group_size:
                dp[node] = child_sum1 + child_sum2 + node_val
            elif min(child_sum1, child_sum2) + node_val <= group_size:
                dp[node] = min(child_sum1, child_sum2) + node_val
                group_count += 1
            else:
                dp[node] = node_val
                group_count += 2
            if group_count > k:
                return False
        return True

    parents = [None] * len(num)
    children_counts = []
    nodes = []
    for node, children in enumerate(links):
        children_count = 0
        for c in children:
            if c != -1:
                parents[c] = node
                children_count += 1
        children_counts.append(children_count)
        if children_count == 0:
            nodes.append(node)
    for i in range(len(num) - 1):
        p = parents[nodes[i]]
        children_counts[p] -= 1
        if children_counts[p] == 0:
            nodes.append(p)

    return algorithm.binary_search(max(num), sum(num) + 1, can_make_group)

• Dependency: teflib.algorithm.binary_search