두 개의 숫자 리스트 nums와 costs가 있다고 가정해 봅시다. 여기서 사용할 수 있는 연산은 nums[i]의 값을 1만큼 증가 또는 감소시킬 때 비용 costs[i]를 지불하는 것입니다. 이 연산은 원하는 만큼 몇 번이든 수행할 수 있으며, 우리의 목표는 nums의 모든 요소를 동일한 값으로 만드는 것입니다. 이때 지불해야 하는 최소 총비용을 구하는 것이 문제입니다.
문제 예시
예를 들어 입력이 다음과 같다고 해봅시다.
- nums = [3, 2, 4]
- costs = [1, 10, 2]
이 경우 정답은 5입니다. 그 이유는 다음과 같습니다.
- 숫자 3을 비용 1로 2로 감소시킵니다.
- 숫자 4를 비용 2씩 두 번 감소시켜 2로 만듭니다.
총비용은 1 + 2 + 2 = 5가 되며, 이보다 적은 비용으로 모든 요소를 같게 만들 수 없습니다.
접근 방법: 볼록 함수와 이진 탐색
핵심 아이디어는 목표값(target)에 대한 총비용 함수가 볼록(convex)하다는 점입니다. 목표값을 낮은 쪽에서 높은 쪽으로 옮기면 총비용은 먼저 감소하다가 최솟값을 지난 후 다시 증가합니다. 따라서 이진 탐색(binary search)을 이용해 비용이 최소가 되는 지점을 효율적으로 찾을 수 있습니다.
풀이 단계
- helper(target) 함수 정의: 특정 목표값 target으로 모든 요소를 맞출 때의 총비용을 계산합니다.
- total := 0 으로 초기화
- enumerate(nums)로 각 인덱스 i와 값 n에 대해 반복
- n이 target과 다르면 total := total + |n − target| × costs[i]
- total 반환
- 메인 로직(이진 탐색):
- low := 0, high := max(nums) 로 범위 설정
- low < high인 동안 반복:
- mid := (low + high) ÷ 2
- helper(mid) < helper(mid + 1)이면 high := mid (최솟값이 왼쪽에 있음)
- 그렇지 않으면 low := mid + 1 (최솟값이 오른쪽에 있음)
- 반복이 끝나면 helper(low) 반환
이 방식의 시간 복잡도는 O(n log m)입니다. 여기서 n은 리스트의 길이, m은 nums의 최댓값입니다. 각 이진 탐색 단계마다 helper 함수가 O(n)으로 실행되기 때문입니다.
구현 예제
class Solution: def solve(self, nums, costs): def helper(target): total = 0 for i, n in enumerate(nums): if target != n: total += abs(n - target) * costs[i] return total low, high = 0, max(nums) while low < high: mid = low + high >> 1 if helper(mid) < helper(mid + 1): high = mid else: low = mid + 1 return helper(low) ob = Solution() nums = [3, 2, 4] costs = [1, 10, 2] print(ob.solve(nums, costs))
입력
[3, 2, 4], [1, 10, 2]
출력
5
정리
각 요소를 변경하는 비용이 서로 다를 때 모든 값을 하나의 값으로 통일하는 문제는, 비용 함수의 볼록성을 활용한 이진 탐색으로 효율적으로 해결할 수 있습니다. 매번 가능한 모든 목표값을 검사하는 대신, 탐색 범위를 절반씩 줄여 나가므로 큰 입력에서도 빠르게 최적해를 찾을 수 있습니다.