[PS] BOJ 2512 / 예산

[PS] BOJ 2512 / 예산
Thumbnail: Photo by Towfiqu barbhuiya (Unsplash)
문제 링크: https://www.acmicpc.net/problem/2512

오랜만에 푸는 매개변수 탐색 문제입니다.

풀이

이분탐색과 매개변수 탐색 돌아보기

1) 이분 탐색은 무엇인가?

이분 탐색은 정렬된 배열 안에서 $O(\log N)$의 시간 복잡도로 배열 안에서 값을 찾는 방법입니다.

2) 결정 문제와 최적화 문제

이분 탐색을 이용해 풀 수 있는 문제 유형 중에, 결정 문제최적화 문제가 있습니다.

  • 결정 문제: Yes / No로 답할 수 있는 문제
    • ex) 배열 A의 최댓값이 3인가?
  • 최적화 문제: 어떤 조건을 만족하는 최댓값/최솟값을 구하는 문제
    • ex) 배열 A의 최댓값을 구하라.

최적화 문제는 결정 문제보다 풀이가 어렵고, 최적화 문제를 풀 수 있다면 그 문제의 결정 문제를 풀 수 있습니다.

3) 최적화 문제의 풀이

최적화 문제는 결정 문제로 바꿔서 풀이할 수 있습니다.

ex) 배열 A의 최댓값을 구하는 문제 (최적화 문제) → 배열 A의 최댓값이 \(K\) 이하인지 판단하는 문제 (결정 문제)

이후, 결정 문제를 \(K = 0, 1, 2, \cdots\) 에 대해 풀면 답을 찾을 수 있습니다.

예를 들어, 배열 A의 최댓값이 10이라 하면, \(K\)에 따라 푼 각각의 결정 문제의 답은 다음과 같습니다.

  • \(K = 0\) → No
  • \(K = 1\) → No
  • ...
  • \(K = 10\) → Yes
  • \(K = 11\) → Yes
  • ...

이를 배열처럼 나타내 보면, NNNNNNNNNYYY.... 꼴입니다.

즉, 배열 A의 최댓값을 구하는 최적화 문제는 결정 문제에서 사용한 임의의 \(K\) 중 결정 문제의 답이 처음으로 Y가 나오는 \(K\)를 찾는 문제로 바꿔 풀 수 있습니다.

정리하면 다음과 같습니다:

  • 조건을 만족하는 최댓값을 탐색
    • YYYYYYNNNN... 에서 가장 마지막으로 나오는 Y의 위치를 찾는 문제
    • 정답이 존재하는 구간 \([S, E]\)와 중간값 \(M\)에 대해, 결정 문제 decision(M)의 해가
      • Y라면 정답은 \([M, E]\)에 존재
      • N이라면 정답은 \([S, M - 1]\)에 존재
    • 이분 탐색 종료 후, \(S\)를 답으로 사용합니다.
  • 조건을 만족하는 최솟값을 탐색
    • NNNNNNYYYY... 에서 가장 처음으로 나오는 Y의 위치를 찾는 문제
    • 정답이 존재하는 구간 \([S, E]\)와 중간값 \(M\)에 대해, 결정 문제 decision(M)의 해가
      • Y라면 정답은 \([S, M]\)에 존재
      • N이라면 정답은 \([M + 1, E]\)에 존재
    • 이분 탐색 종료 후, \(E\)를 답으로 사용합니다.

이처럼 '최적화 문제를 결정 문제로 바꾸어 푸는 알고리즘' 을 매개변수 탐색(Parametric Search) 라고 합니다.

문제 풀이

예산 상한액을 $M$이라 할 때 $M$의 최댓값을 찾는 문제입니다.

예산 상한액 $M$을 임의로 결정했을 때, 이 상한액을 토대로 각 부서에 예산을 할당한 뒤 할당된 예산의 합을 $S$라 합시다. 만약 $S$가 국가 예산의 총액을 넘어설 경우, 현재 상한액 $M$은 불가능한 경우입니다 (N). 그렇지 않은 경우는 예산 상한액으로 사용할 수 있습니다 (Y).

각 부서별 예산은 $1 \leq 100,000$이므로, 예산 범위 안의 모든 수를 상한액으로 사용할 때 그 결과를 YYYY...YNNNN 꼴로 변환할 수 있습니다. 다시 말해, 상한액 $M$을 찾는 매개변수 탐색 문제로 생각할 수 있습니다.

처음 매개변수 탐색 범위는 예산으로 주어질 수 있는 범위와 같습니다. ($l = 1, r = 100,000$)
이후, $M = (l + r) / 2$로 계산한 뒤 $M$이 예산 상한액일 때 각 부서에 할당될 예산의 합을 계산합니다.

예산 합 $S$가 국가 예산의 총액보다 같거나 작다면 Y이므로 $l$을 $M + 1$로 변경합니다. 그렇지 않다면, N이므로 $r$을 $M - 1$로 변경합니다.

이후, 매개변수 탐색이 종료될 때 $r$을 예산의 상한액으로 사용하면 됩니다.

전체 코드

input = open(0).readline

N = int(input())
budgets = list(map(int, input().split()))
total_budget = int(input())

max_budget = 0
total = 0
for b in budgets:
    total += b
    max_budget = max(max_budget, b)

if total <= total_budget:
    print(max_budget)
else:
    left = 0  # 기관별 예산 요구값의 최솟값조차 상한액보다 클 수 있다.
    right = 100000  # 어차피 상한액이 예산 최대값보다 클 필요는 없으니까
    while (left <= right):
        mid = (left + right) // 2

        # mid를 상한액으로 정하고 문제를 풀 때 답 구하기
        total = 0
        for i in range(N):
            total += min(budgets[i], mid)
        
        if total <= total_budget: # 답이 Y: 정답은 [mid, right] 구간에 존재
            left = mid + 1
        else: # 답이 N: 정답은 [left, mid-1] 구간에 존재
            right = mid - 1
    print(right)

solution.py

시행착오

이분 탐색의 범위 설정 문제

처음엔 예산 배열을 정렬하고 사용했습니다. 이는 매개변수 탐색의 결정 문제인 "상한액 $M$에 대해 할당된 예산의 합 $S$가 국가 예산 총액보다 작거나 같은가?"를 풀 때, $S$를 더 빠르게 계산하고 또 기관들이 요구한 예산의 최댓값을 알기 위해서 였습니다.

예산 배열을 정렬해 얻은 예산 최댓값을 $r$로, $l$은 0으로 초기값을 정하고 매개변수 탐색을 진행하면서, $M$(mid)은 다음과 같이 조정했습니다.

  • $S$ <= 국가 예산 총액
    • $l = M$ (답은 $[M, r]$에 있다)
  • 그렇지 않으면
    • $r = M - 1$ (답은 $[l, M - 1]$에 있다)
N = int(input())
budgets = sorted(map(int, input().split()))
total_budget = int(input())

if sum(budgets) <= total_budget:
    print(budgets[N - 1])
else:
    left = 0  # 기관별 예산 요구값의 최솟값조차 상한액보다 클 수 있다.
    right = budgets[N - 1]  # 어차피 상한액이 예산 최대값보다 클 필요는 없으니까
    while (left <= right):
        mid = (left + right) // 2

        # mid를 상한액으로 정하고 문제를 풀 때 답 구하기
        total = 0
        for i in range(N):
            if budgets[i] >= mid:
                total += mid * (N - i) # 이 이후의 예산은 모두 상한액으로 배정된다.
                break
            else:
              total += budgets[i]
        
        if total <= total_budget: # 답이 Y: 정답은 [mid, right] 구간에 존재
            left = mid
        else: # 답이 N: 정답은 [left, mid-1] 구간에 존재
            right = mid - 1
    print(right)

초기 코드

하지만 이렇게 구현한 이분 탐색은 무한 루프에 빠지게 됩니다. 이는 이분 탐색의 종료 조건 ($l \leq r$)에 의해 $l = r$인 경우는 while문이 종료되지 않는데, $l = M$으로 설정하면 다음 반복에도 다시 $l = M$이 되풀이되며 무한히 반복합니다.

따라서, $l = M + 1$로 조정되어야 합니다.

또, $r$을 예산의 최댓값으로 정할 경우 몇몇 경우에서 오답을 구하게 됩니다. 각 부서가 요구할 수 있는 예산의 최댓값인 100,000을 $r$로 설정하면 정답을 구할 수 있습니다.

배열을 굳이 정렬할 필요가 있을까?

앞서 이분 탐색의 범위를 고치며, 굳이 예산의 최댓값이 필요하지 않게 되었습니다.
또, 예산 배열을 정렬해서 이후 결정 문제의 풀이 시간을 단축하는 것 보다 정렬 자체의 시간이 더 오래 걸릴 수 있음을 고려해 정렬을 하지 않도록 변경했습니다.

정렬을 하지 않으니, 결정 문제에서 $S$를 구하는 코드는 각 부서의 예산과 상한액 중 작은 값을 더하게끔 구현해주면 됩니다.

# mid를 상한액으로 정하고 문제를 풀 때 답 구하기
total = 0
for i in range(N):
    total += min(budgets[i], mid)

개선된 결정 문제 풀이

또, 국가 예산 총액이 각 부서별로 요구한 예산의 합보다 더 크다면 요구한대로 예산을 줄 수 있으니 이분탐색을 진행하지 않고 요구한 예산의 최댓값을 바로 출력했습니다. 배열을 정렬했을 때는 항상 마지막 원소가 최댓값이었으나, 이제 배열을 정렬하지 않으니 예산의 합을 계산하면서 최댓값을 같이 찾았습니다.

max_budget = 0
total = 0
for b in budgets:
    total += b
    max_budget = max(max_budget, b)

if total <= total_budget:
    print(max_budget)
else:
    ...

예외 처리