문제 소개
여러 개의 반이 있고, 각 반은 classes[i] = [pass_i, total_i] 형태로 표현됩니다. 여기서 pass_i는 i번째 반에서 시험에 합격한 학생 수, total_i는 해당 반의 전체 학생 수를 의미합니다. 또한 추가로 배정할 수 있는 우수한 학생 수 extra가 주어지는데, 이 학생들은 어느 반에 배정되든 반드시 시험에 합격하는 것이 보장됩니다.
우리의 목표는 이 추가 학생들을 각 반에 배치하여, 모든 반의 평균 통과율을 최대화하는 것입니다. 통과율(pass ratio)은 해당 반에서 합격한 학생 수를 전체 학생 수로 나눈 값이며, 평균 통과율은 모든 반의 통과율을 합한 뒤 반의 개수로 나눈 값입니다.
예시
예를 들어 classes = [[2,3],[4,6],[3,3]], extra = 3이라면 출력은 0.83809가 됩니다. 첫 번째 반에 학생 2명, 두 번째 반에 학생 1명을 배정하는 것이 가장 유리하기 때문입니다. 이 경우 평균은 (4/5 + 5/7 + 3/3) ÷ 3 = 0.83809로 계산됩니다.
접근 방법
이 문제의 핵심은 그리디(Greedy) 기법입니다. 추가 학생 한 명을 배정할 때마다 통과율이 가장 크게 향상되는 반을 선택하면 됩니다. 이 선택을 효율적으로 수행하기 위해 최대 힙(max-heap)을 사용합니다. Python의 heapq는 최소 힙만 지원하므로, 통과율 변화량에 음수를 취해 저장하는 방식으로 구현합니다.
구체적인 단계는 다음과 같습니다.
- 각 반 (a, b)에 대해
(a/b − (a+1)/(b+1), a, b)형태의 튜플 목록 h를 만듭니다. 이 값은 학생 1명을 추가했을 때의 통과율 변화량(감소분)을 나타냅니다. - h를 힙으로 변환(
heapify)합니다. - extra가 0이 될 때까지 다음 과정을 반복합니다.
- 힙에서 최상위 요소(통과율 향상폭이 가장 큰 반)를 꺼냅니다.
- (a, b)를 (a+1, b+1)로 갱신합니다.
- 갱신된 값과 새로운 변화량을 계산해 힙에 다시 삽입합니다.
- extra를 1 감소시킵니다.
- 모든 배정이 끝나면 힙에 남은 반들의 통과율 평균을 반환합니다.
이 방식의 시간 복잡도는 O((n + extra) log n)으로, 매번 모든 반을 탐색하는 O(n × extra) 방식보다 훨씬 효율적입니다.
구현 예제
다음 코드를 통해 더 잘 이해해 보겠습니다.
import heapq
def solve(classes, extra):
h = [(a / b - (a + 1) / (b + 1), a, b) for a, b in classes]
heapq.heapify(h)
while extra:
v, a, b = heapq.heappop(h)
a, b = a + 1, b + 1
heapq.heappush(h, (-(a + 1) / (b + 1) + a / b, a, b))
extra -= 1
return sum(a / b for v, a, b in h) / len(h)
classes = [[2,3],[4,6],[3,3]]
extra = 3
print(solve(classes, extra))
입력
[[2,3],[4,6],[3,3]], 3
출력
0.8380952380952381
마무리
이처럼 힙 자료구조를 활용하면 '매 순간 가장 이득이 되는 선택'을 빠르게 찾아낼 수 있습니다. 추가 학생을 한 명씩 배정할 때마다 통과율 개선 폭이 가장 큰 반을 O(log n) 시간에 선택함으로써, 전체 평균 통과율을 효율적으로 최대화할 수 있습니다.