백준 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)\ast a[m]$ 에

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

 

 

 

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

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

 

시간복잡도는 $O(RlogR)$이다. (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 ########