[프로그래머스] 시험장 나누기 (2021 카카오 채용연계형 인턴십)

https://school.programmers.co.kr/learn/courses/30/lessons/81305#

프로그래머스

 

이진트리를 k개 그룹으로 분할할 때, 각 그룹에 있는 노드가 가진 가중치의 합의 최댓값이 최소가 되도록 하는 문제

최댓값이 최소가 되어야 한다는 점에서 이분탐색 (매개변수 탐색) 까지는 어렵지 않게 떠올릴 수 있었으나, 이진트리를 k개로 분할하는 구현에서 막혀서 어려웠다.

 

처음에는 각 트리에 대해서 누적합을 구하고, 누적합 정보를 기반으로 쪼개는 풀이를 떠올렸다.

하지만 이 풀이가 27점 맞고 틀려서 디버깅하다가 반례를 못찾아서 클로드에 물어봤더니 클로드는 정답 풀이를 잘 짜주었다. (무료버전)

GPT는 문제는 이해했으나 풀이를 내는데에는 실패했는데... 클로드가 이젠 더 똑똑해진 것 같다.

 

클로드의 풀이와 내 풀이를 비교하면서 클로드로부터 반례도 받았는데, GPT와 다르게 반례도 제대로 만들어줘서 도움을 많이 받았다.

내 풀이는 탑-다운으로 누적합 데이터를 이용해서 위에서부터 내려가면서 그룹을 나눴는데, 이렇게 나누니까 

 

      10

     / 

  10

  /

10

   \

    10

 

를 2개로 분할할 때 20이 아니라 30을 가능한 최댓값의 최소로 잡는 문제가 생겼다.

20을 기준으로 잡았을 때

루트 10 / 그 아래의 10 / 나머지 10 10 

이렇게 3개로 분할했기 때문이다.

 

위에서부터 내려오면 루트 기준 누적합이 40이므로 20보다 커서 쪼개기로 판단한다.

이 상태에서 일단 바로 쪼개버리기 때문에 벌써 오답이다.

그래서 쪼개는 판단은 최대한 미뤄서 나중에 판단해야 한다.

 

import sys
sys.setrecursionlimit(10 ** 5)


def solution(k, num, links):
    acc = [0 for _ in range(len(num))]

    def get_acc(i):
        if acc[i] > 0:
            return acc[i]
        if i == -1:
            return 0

        left, right = links[i]
        get_left, get_right = 0, 0
        if left != -1:
            get_left = get_acc(left)
        if right != -1:
            get_right = get_acc(right)

        acc[i] = num[i] + get_left + get_right
        return acc[i]

    #  트리의 루트 찾기
    root = -1
    s = sum(num)
    for i in range(len(num)):
        a = get_acc(i)
        if a == s:
            root = i

    def count_cuttings(node, limit):
        if node == -1:
            return 0, 0

        left, right = links[node]
        left_count, left_sum = count_cuttings(left, limit)
        right_count, right_sum = count_cuttings(right, limit)

        # 현재 서브트리 전체를 그룹으로 만들 수 있는 경우
        if num[node] + left_sum + right_sum <= limit:
            print(node, 0, num[node] + left_sum + right_sum)
            return right_count + left_count, num[node] + left_sum + right_sum

        # 좌/우 서브트리 중 하나만 그룹에 포함시킬 수 있는 경우
        # 더 작은 사이즈의 트리를 합쳐서 그룹에 넣고, 큰 그룹은 분리
        if num[node] + min(left_sum, right_sum) <= limit:
            return right_count + left_count + 1, num[node] + min(left_sum, right_sum)

        # 좌/우 서브트리 모두 그룹에 포함시킬 수 없는 경우
        return right_count + left_count + 2, num[node]


    def possible(v):
        count, _ = count_cuttings(root, v)
        return count < k

    start, end = max(num) - 1, s
    while start + 1 < end:
        mid = (start + end) // 2
        if possible(mid):
            end = mid
        else:
            start = mid

    answer = end
    return answer

 

프로그래머스 AI가 추천해줘서 풀어봤는데 알고보니 Lv.5 문제였다.

이 문제를 실제 시험에서 시간 내에 푼 사람이 있다니 대단하다는 생각이 들었다.

알고리즘도 열심히 해야겠다..