結果

問題 No.738 平らな農地
ユーザー norioc
提出日時 2025-03-12 03:54:57
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 4,527 bytes
コンパイル時間 399 ms
コンパイル使用メモリ 82,320 KB
実行使用メモリ 99,392 KB
最終ジャッジ日時 2025-03-12 03:55:07
合計ジャッジ時間 9,827 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 5
other AC * 10 TLE * 1 -- * 76
権限があれば一括ダウンロードができます

ソースコード

diff # 

import random


class RBST:
    class Node:
        def __init__(self, val):
            self.val = val
            self.left = None
            self.right = None
            self.size = 1
            self.sum = val

        def update(self):
            pass

        def push(self):
            pass

    def __init__(self, node: Node = None):
        self.root = node

    def __len__(self):
        return RBST._size(self.root)

    def __iter__(self):
        for i in range(len(self)):
            yield self.get(i)

    @staticmethod
    def _sum(node: Node) -> int:
        if node: return node.sum
        return 0

    @staticmethod
    def _size(node: Node) -> int:
        if node is None: return 0
        return node.size

    @staticmethod
    def _update(node: Node) -> Node:
        node.size = RBST._size(node.left) + RBST._size(node.right) + 1
        node.sum = RBST._sum(node.left) + RBST._sum(node.right) + node.val
        node.update()
        return node

    @staticmethod
    def _push(node: Node):
        if node is None: return
        node.push()

    @staticmethod
    def _lower_bound(node: Node, val) -> int:
        RBST._push(node)
        if node is None: return 0
        if val <= node.val:
            return RBST._lower_bound(node.left, val)
        else:
            return RBST._size(node.left) + RBST._lower_bound(node.right, val) + 1

    def lower_bound(self, val) -> int:
        return RBST._lower_bound(self.root, val)

    @staticmethod
    def _upper_bound(node: Node, val) -> int:
        RBST._push(node)
        if node is None: return 0
        if val >= node.val:
            return RBST._size(node.left) + RBST._upper_bound(node.right, val) + 1
        else:
            return RBST._upper_bound(node.left, val)

    def upper_bound(self, val) -> int:
        return RBST._upper_bound(self.root, val)

    def count(self, val) -> int:
        return self.upper_bound(val) - self.lower_bound(val)

    @staticmethod
    def _get(node: Node, k: int):
        RBST._push(node)
        if node is None: return -1
        if k == RBST._size(node.left): return node.val
        if k < RBST._size(node.left):
            return RBST._get(node.left, k)
        else:
            return RBST._get(node.right, k - RBST._size(node.left) - 1)

    def get(self, k):
        return RBST._get(self.root, k)

    @staticmethod
    def _merge(left: Node, right: Node) -> Node:
        RBST._push(left)
        RBST._push(right)
        if left is None or right is None:
            if left: return left
            else: return right

        r = random.randint(1, 0xfffffff)
        if r % (left.size + right.size) < left.size:
            left.right = RBST._merge(left.right, right)
            return RBST._update(left)
        else:
            right.left = RBST._merge(left, right.left)
            return RBST._update(right)

    def merge(self, add: 'RBST'):
        self.root = RBST._merge(self.root, add.root)

    @staticmethod
    def _split(node: Node, k: int) -> tuple[Node, Node]:
        # [0, k), [k, n)
        RBST._push(node)
        if node is None: return node, node
        if k <= RBST._size(node.left):
            a, b = RBST._split(node.left, k)
            node.left = b
            return a, RBST._update(node)
        else:
            a, b = RBST._split(node.right, k - RBST._size(node.left) - 1)
            node.right = a
            return RBST._update(node), b

    def split(self, k: int):
        a, b = RBST._split(self.root, k)
        self.root = a
        return RBST(b)

    def insert(self, val):
        a, b = RBST._split(self.root, self.lower_bound(val))
        self.root = RBST._merge(RBST._merge(a, RBST.Node(val)), b)

    def remove(self, val):
        if self.count(val) == 0: return
        a, b = RBST._split(self.root, self.lower_bound(val))
        _, r = RBST._split(b, 1)
        self.root = RBST._merge(a, r)

    def sum(self):
        return RBST._sum(self.root)


INF = 1 << 60
N, K = map(int, input().split())
A = list(map(int, input().split()))


def calc_cost():
    if K == 1: return 0
    if K == 2:
        return bst.get(1) - bst.get(0)

    mid = bst.split(K // 2)
    rt = mid.split(1)

    m = mid.get(0)
    res = m * len(bst) - bst.sum()
    res += rt.sum() - m * len(rt)

    bst.merge(mid)
    bst.merge(rt)
    return res


ans = INF
bst = RBST()
for i, a in enumerate(A):
    bst.insert(a)
    if len(bst) > K:
        bst.remove(A[i-K])

    if len(bst) == K:
        ans = min(ans, calc_cost())

print(ans)
0