Python으로 가방 속 공의 최소 페널티 찾기

문제 설명

정수 배열 nums가 주어지며, 각 원소 nums[i]는 i번째 가방에 들어 있는 공의 개수를 나타냅니다. 또한 연산을 수행할 수 있는 최대 횟수를 의미하는 값 mx가 함께 주어집니다. 우리가 사용할 수 있는 연산은 다음과 같습니다.

공이 든 가방 하나를 선택해, 각각 최소 한 개 이상의 공을 포함하도록 두 개의 새로운 가방으로 나눕니다.
페널티(penalty)란 모든 가방 중 가장 많은 공을 담고 있는 가방의 공 개수를 의미합니다.

연산을 최대 mx번까지 수행한 뒤 얻을 수 있는 최소 페널티를 구하는 것이 이 문제의 목표입니다.

예를 들어 입력이 nums = [4,8,16,4], mx = 4라고 하면 결과는 4입니다. 처음 가방 상태가 [4,8,16,4]일 때, 공 16개가 든 가방을 나누면 [4,8,8,8,4]가 됩니다. 이후 공 8개가 든 가방들을 차례로 반씩 나누어 [4,4,4,8,8,4] → [4,4,4,4,4,8,4] → 마지막으로 [4,4,4,4,4,4,4,4] 순서로 만들 수 있습니다. 최종적으로 가장 큰 가방에도 공이 4개뿐이므로, 최소 페널티는 4입니다.

접근 방법: 이진 탐색

이 문제는 이진 탐색(binary search)으로 효율적으로 해결할 수 있습니다. 핵심 아이디어는 “페널티가 정확히 t 이하가 되도록 만드는 것이 가능한가?”라는 질문을 임의의 값 t에 대해 빠르게 판별하고, 조건을 만족하는 t의 최솟값을 탐색하는 것입니다.

판별 함수(helper)의 동작 원리

n개의 공이 든 가방을 한 가방당 최대 t개의 공만 남도록 나누려면 필요한 분할 횟수는 (n - 1) // t입니다. 예를 들어 n = 16, t = 8이면 (16 − 1) // 8 = 1회, n = 16, t = 4이면 15 // 4 = 3회가 필요합니다. 따라서 모든 가방에 필요한 분할 횟수의 합이 mx 이하라면 페널티 t를 달성할 수 있습니다.

helper(target, mx) 함수를 정의합니다.
target이 0이면 mx + 1을 반환합니다(나눌 수 없는 경우를 처리).
count를 0으로 초기화한 뒤, 각 가방에 대해 count += (num - 1) // target을 누적합니다.
count <= mx 여부를 반환합니다. 참이면 mx번 이내의 연산으로 페널티를 target 이하로 만들 수 있다는 뜻입니다.

탐색 범위 정하기

하한(left): 전체 공은 최대 len(nums) + mx개의 가방에 담길 수 있으므로, 페널티는 평균값 sum(nums) // (len(nums) + mx)보다 작아질 수 없습니다. 다만 최솟값이 1이 되도록 1과 비교해 더 큰 값을 사용합니다.
상한(right): 연산을 전혀 수행하지 않아도 되므로, 초기 페널티의 최댓값은 max(nums)입니다.

이 범위에서 이진 탐색을 진행하여, 조건을 만족하는 가장 작은 값을 찾아 반환하면 됩니다.

구현 예제

다음 코드를 통해 더 쉽게 이해할 수 있습니다.

def helper(target, mx):
    if target == 0:
        return mx + 1
    count = 0
    for num in nums:
        count += (num - 1) // target
    return count <= mx

def solve(nums, mx):
    left, right = max(sum(nums) // (len(nums) + mx), 1), max(nums)
    while left < right:
        mid = (left + right) // 2
        if helper(mid, mx):
            right = mid
        else:
            left = mid + 1
    return left

nums = [4,8,16,4]
mx = 4
print(solve(nums, mx))

입력

[4,8,16,4], 4

출력

시간 복잡도 분석

helper 함수는 배열을 한 번 순회하므로 O(n)이고, 이진 탐색은 최대 log(max(nums))번 반복됩니다. 따라서 전체 시간 복잡도는 O(n · log(max(nums)))이며, 추가 공간은 상수 수준(O(1))으로 매우 효율적입니다. 완전 탐색으로 모든 경우를 시도하는 것과 비교하면 입력 크기가 커져도 안정적인 성능을 보장한다는 점이 큰 장점입니다.

Python으로 가방 속 공의 최소 페널티 찾기 – 이진 탐색 완전 정복