백준 2805번의 시간 절감을 위한 풀이

백준 2805번 문제를 보면 대부분 이분 탐색을 사용하여 문제를 풀게 되는데 Python3로 돌리면 시간 초과가 나서 PyPy3로만 성공할 수 있다는 것이 조금 마음에 걸렸다.

일단 이분 탐색을 이용한 정석적 풀이를 먼저보면 아래와 같다.

n, m = map(int, input().split())
tree = list(map(int, input().split()))
start, end = 1, max(tree) 

while start <= end: 
    mid = (start+end) // 2
    
    total = 0 
    for i in tree:
        if i >= mid:
            total += i - mid
    
    if total >= m:
        start = mid + 1
    else:
        end = mid - 1
print(end)

코드는 심플해서 매우 좋다.

조금 더 빠른 계산을 위해 이렇게 생각해보자.

1. 나무를 크기 순으로 쭉 나열한다.
2. 절단기를 제일 높은 나무 꼭대기에서부터 1씩 높이를 내린다.
3. 두 번째로 높은 나무 꼭대기에 다다를 때까지는 절단기를 1 내릴 때마다 나무의 양도 1씩 증가한다.
4. 두 번째로 높은 나무의 높이보다 절단기가 낮아지면 절단기를 1 내릴 때마다 나무의 양은 2씩 증가한다. 이러한 증가는 절단기가 세 번째로 높은 나무 꼭대기에 닿을 때까지 계속된다.
5. 절단기를 1씩 내리다가 나무의 양이 목표량에 도달하면 절단기 높이를 측정한다.

이를 코드로 표현하자면

import sys
n, goal = map(int, input().split())
tree = list(map(int, sys.stdin.readline().split())) # input 받아오기

tree.sort(reverse=True) # 나무를 높이순으로 정렬
tree.append(0) # 아래의 for문에서 인덱스 오류를 막기 위해 추가
total = 0 # 수확한 나무의 양
height = 0 # 절단기를 1씩 내린 총 횟수
for i in range(n):
    total += (tree[i]-tree[i+1])*(i+1) 
    # i번째로 높은 나무에서 i+1번째로 높은 나무 꼭대기까지 절단기를 내릴 떄는 i+1씩 나무량이 증가
    height += (tree[i]-tree[i+1]) # 1씩 내리면 계산량이 많으므로 나무 꼭대기와 꼭대기 사이를 한번에 계산
    if total >= goal: # 목표량을 넘어버렸으면 절단기 높이를 올려서 나무량을 낮춤
        height -= (total-h)//(i+1) # 절단기를 얼마나 올려야하는지 계산
        break
print(tree[0]-height)

이와 같이 된다. Python3로도 준수한 속도로 통과가 가능하다!