結果

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

ソースコード

diff #
raw source code 

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()))

if K == 1:
    print(0)
    exit()


def make_rbst(xs):
    bst = RBST()
    for x in xs:
        bst.insert(x)

    m = (K + 1) // 2
    rt = bst.split(m)
    return bst, rt


def calc_cost() -> int:
    m = (K + 1) // 2

    mid = lt.get(m-1)
    lsum = lt.sum() - mid
    rsum = rt.sum()

    res = mid * (len(lt)-1) - lsum
    res += rsum - mid * len(rt)
    return res


def replace_and_relax(old, new):
    m = (K + 1) // 2
    lx = lt.get(m-1)
    if old <= lx:
        lt.remove(old)
        lt.insert(new)
    else:
        rt.remove(old)
        rt.insert(new)


    lx = lt.get(m-1)
    rx = rt.get(0)

    if lx > rx:
        lt.remove(lx)
        rt.remove(rx)
        lt.insert(rx)
        rt.insert(lx)


lt, rt = make_rbst(A[:K])
ans = calc_cost()
for i in range(K, N):
    replace_and_relax(A[i-K], A[i])
    ans = min(ans, calc_cost())

print(ans)
0