백준 5821 - 쌀 창고 [Python]

https://www.acmicpc.net/problem/5821

5821번: 쌀 창고

 

- 이분 탐색

- 누적 합

 

논의 개수가 홀수일 때, 가운데에 있는 논에 창고를 세울 때 최소비용이 된다.

논의 개수가 짝수일 때는, 가운데 두 논의 좌표 a, b에 대해 a <= k <= b를 만족하는 k에 창고를 지으면 최소비용이 된다. 

 

간단하게, 논의 개수가 짝수일 때 왼쪽의 논을 고른다고 하면,

l 부터 r 까지의 논에서의 최소 비용을 만드는 논의 위치는 \( \lfloor \frac{l + r}{2} \rfloor \)이다.

 

그러한 위치를 \(m\)이라고 하자.

 

\(p[i]\)를 "i번째 논 까지의 좌표의 합" 이라고 하면 \(l\) ~ \(r\)의 비용은,

중앙 ~ 오른쪽의 비용인 \( p[r] - p[m] - (r - m) * a[m] \) 에

왼쪽 ~ 중앙의 비용인 \( (m - l) * a[m] - (p[m - 1] - p[l - 1]) \) 를 더한 값이 된다.

 

 

 

이제, 답이 \( X \)일 때, 이 \(X\)에 대해 이분탐색을 진행할 것이다.

\(X\)개의 논을 포함하는 창고를 지었을 때의 비용이, 예산 \(B\)보다 작거나 같다면 답을 갱신하면 된다.

 

시간복잡도는 \(O(R log R)\)이다. (R은 논의 개수)

 

import sys, math
import heapq
 
from collections import deque
input = sys.stdin.readline
 
#input
def ip(): return int(input())
def sp(): return str(input().rstrip())
 
def mip(): return map(int, input().split())
def msp(): return map(str, input().split().rstrip())
 
def lmip(): return list(map(int, input().split()))
def lmsp(): return list(map(str, input().split().rstrip()))
 
############### Main! ###############
 
MAX = 100003
p = [0 for _ in range(MAX)]
n, k, b = mip()
a = []
for i in range(n):
    a.append(ip())
 
for i in range(n):
    p[i] = p[i - 1] + a[i]
 
l = 1
r = n
ans = 0
while l <= r:
    m = (l + r) // 2
    c = 0
    for i in range(n - m + 1):
        x = i
        y = i + m - 1
        mid = (x + y) // 2
        z = p[y] - p[mid] - (y - mid) * a[mid] + (mid - x) * a[mid] - (p[mid - 1] - p[x - 1])
        if z <= b:
            c = 1
            break
    if c:
        l = m + 1
        ans = max(ans, m)
    else:
        r = m - 1
 
print(ans)
 
######## Priest W_NotFoundGD ########