数列を上位半分と下位半分に分けて中央値を管理するやつ
(expansion/misc/MadianHeap.py)

• View this file on GitHub
• Last update: 2023-07-16 15:47:36+09:00

概要

中央値を管理する用のものです．内部で hashDict を使っているので，hack 対策に rnd() を一様に足してあげると良いかもしれません

使い方

• __init__(A=None) := Aに要素を指定すると初期の配列を作ってくれます
• __len__() :=現在の配列長を返します．
• push(x) := xを追加します．
• remove(x) := xを削除します．存在しない場合未定義です
• get_med(self)
  • 数列長が奇数 (2n+1) の場合: 0-indexed で [A[n], A[n]] を返します．
  • 数列長が偶数 (2n) の場合: 0-indexed で [A[n - 1], A[n]] を返します．

• abs_sum() := 「数列Aが与えられる．$\sum_{a \in A}

  a-x

  $ の最小値を求めろ」の答えを返します．x の値は get_med() の返り値がそのまま下限と上限です．

Code

import heapq


class DeletableHeapq:
    def __init__(self, lst=None, reverse=False):
        if reverse:
            self.pm = -1
        else:
            self.pm = 1
        if lst is None:
            self.hq = []
        else:
            self.hq = [self.pm * x for x in lst]

        heapq.heapify(self.hq)
        self.tot = sum(self.hq) * self.pm
        self.cnt = {}
        for x in self.hq:
            self.cnt[x] = self.cnt.get(x, 0) + 1
        self.length = len(self.hq)

    def __bool__(self):
        return bool(self.hq)

    def __len__(self):
        return self.length

    @property
    def sum(self):
        return self.tot

    def top(self):
        if self.hq:
            return self.hq[0] * self.pm
        else:
            return None

    def __getitem__(self, i):
        # 先頭要素にアクセスしたいときのみ
        assert i == 0
        return self.top()

    def push(self, x):
        self.length += 1
        self.cnt[x * self.pm] = self.cnt.get(x * self.pm, 0) + 1
        self.tot += x
        heapq.heappush(self.hq, x * self.pm)

    def pop(self):
        if not self.hq:
            return None
        self.length -= 1
        x = heapq.heappop(self.hq)
        self.cnt[x] -= 1
        self.tot -= x * self.pm
        self._delete()
        return x * self.pm

    def remove(self, x):
        if self.cnt.get(x * self.pm, 0) == 0:
            return False
        self.length -= 1
        self.cnt[x * self.pm] -= 1
        self.tot -= x
        self._delete()
        return True

    def _delete(self):
        while self.hq and self.cnt.get(self.hq[0], 0) == 0:
            heapq.heappop(self.hq)


class MedianHeap:
    def __init__(self, A=None):
        self.upp = DeletableHeapq()
        self.low = DeletableHeapq(reverse=True)

        if A is not None:
            for a in A:
                self.push(a)

    def __len__(self):
        return len(self.upp) + len(self.low)

    def push(self, x):
        if len(self.upp) == len(self.low) + 1:
            self.upp.push(x)
            self.low.push(self.upp.pop())
        else:
            self.low.push(x)
            self.upp.push(self.low.pop())

    def get_med(self):
        """
        A = [1, 2, 3] -> [2, 2] を返す
        A = [1, 2, 3, 4] -> [2, 3] を返す
        """
        if len(self.upp) == len(self.low):
            return [self.low[0], self.upp[0]]
        else:
            return [self.upp[0], self.upp[0]]

    def remove(self, x):
        if x <= self.low[0]:
            self.low.remove(x)
            if len(self.low) + 1 < len(self.upp):
                self.low.push(self.upp.pop())
        else:
            self.upp.remove(x)
            if len(self.upp) < len(self.low):
                self.upp.push(self.low.pop())

    def abs_sum(self):
        if len(self) == 0:
            return 0

        med = self.upp[0]
        ret = med * len(self.low) - self.low.tot
        ret += self.upp.tot - med * len(self.upp)
        return ret

Traceback (most recent call last):
  File "/opt/hostedtoolcache/Python/3.11.4/x64/lib/python3.11/site-packages/onlinejudge_verify/documentation/build.py", line 81, in _render_source_code_stat
    bundled_code = language.bundle(
                   ^^^^^^^^^^^^^^^^
  File "/opt/hostedtoolcache/Python/3.11.4/x64/lib/python3.11/site-packages/onlinejudge_verify/languages/python.py", line 108, in bundle
    raise NotImplementedError
NotImplementedError

Back to top page