結果

問題 No.738 平らな農地
ユーザー lam6er
提出日時 2025-03-20 20:21:40
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 723 ms / 2,000 ms
コード長 3,005 bytes
コンパイル時間 345 ms
コンパイル使用メモリ 82,464 KB
実行使用メモリ 103,992 KB
最終ジャッジ日時 2025-03-20 20:24:04
合計ジャッジ時間 30,348 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 5
other AC * 87
権限があれば一括ダウンロードができます

ソースコード

diff # 

import heapq
from collections import defaultdict

class SlidingWindowMedian:
    def __init__(self):
        self.small = []  # max-heap (stores negative values)
        self.large = []  # min-heap
        self.small_size = 0
        self.large_size = 0
        self.small_sum = 0
        self.large_sum = 0
        self.deleted = defaultdict(int)
    
    def add_num(self, num):
        if not self.small or num <= -self.small[0]:
            heapq.heappush(self.small, -num)
            self.small_size += 1
            self.small_sum += num
        else:
            heapq.heappush(self.large, num)
            self.large_size += 1
            self.large_sum += num
        self._rebalance()
    
    def _clean(self):
        while self.small and self.deleted.get(-self.small[0], 0) > 0:
            val = -heapq.heappop(self.small)
            self.deleted[val] -= 1
            if self.deleted[val] == 0:
                del self.deleted[val]
        while self.large and self.deleted.get(self.large[0], 0) > 0:
            val = heapq.heappop(self.large)
            self.deleted[val] -= 1
            if self.deleted[val] == 0:
                del self.deleted[val]
    
    def _rebalance(self):
        self._clean()
        while self.small_size - self.large_size > 1:
            val = -heapq.heappop(self.small)
            self.small_size -= 1
            self.small_sum -= val
            heapq.heappush(self.large, val)
            self.large_size += 1
            self.large_sum += val
            self._clean()
        while self.large_size > self.small_size:
            val = heapq.heappop(self.large)
            self.large_size -= 1
            self.large_sum -= val
            heapq.heappush(self.small, -val)
            self.small_size += 1
            self.small_sum += val
            self._clean()
    
    def remove_num(self, num):
        self.deleted[num] += 1
        if num <= -self.small[0]:
            self.small_size -= 1
            self.small_sum -= num
        else:
            self.large_size -= 1
            self.large_sum -= num
        self._rebalance()
    
    def get_cost(self):
        self._clean()
        if self.small_size + self.large_size == 0:
            return 0
        m = -self.small[0] if self.small_size >= self.large_size else self.large[0]
        cost = (m * self.small_size - self.small_sum) + (self.large_sum - m * self.large_size)
        return cost

def main():
    import sys
    input = sys.stdin.read
    data = input().split()
    n = int(data[0])
    k = int(data[1])
    a = list(map(int, data[2:2+n]))
    
    if k == 0:
        print(0)
        return
    
    swm = SlidingWindowMedian()
    min_cost = float('inf')
    
    for i in range(n):
        swm.add_num(a[i])
        if i >= k - 1:
            current_cost = swm.get_cost()
            if current_cost < min_cost:
                min_cost = current_cost
            swm.remove_num(a[i - k + 1])
    
    print(min_cost)

if __name__ == '__main__':
    main()
0