公式

D - 中央値チェック / Median Check 解説 by admin 印刷ページ

gpt-5.3-codex

概要

クエリで多重集合に値を追加・削除しながら、削除時に「その直前の中央値と削除値が一致した回数」を数える問題です。
各時点の中央値を高速に求めるために、座標圧縮 + Fenwick Tree(BIT) を使って処理します。

考察

まず、削除クエリ - x のたびに「現在の中央値」を求める必要があります。
中央値は、要素数を \(n\) とすると昇順で \(\lceil n/2 \rceil\) 番目、つまりコードでは
\(k = (n+1)//2\) 番目の値です。

素朴解が厳しい理由

例えば毎回、 - 多重集合をソートし直す - またはリストに挿入・削除して中央値位置を見る

とすると、1 クエリあたり \(O(n)\) 以上かかり、\(Q \le 2\times10^5\) では間に合いません(TLE)。

重要な気づき

必要なのは「値そのものの大小関係」と「各値が何個あるか」です。
そこで、

  1. すべての \(X_i\) を事前に集めてソート・重複除去(座標圧縮)
  2. 圧縮後の添字に対して Fenwick Tree で個数管理
  3. Fenwick Tree の累積和を使って「\(k\) 番目の要素」を二分探索的に求める

とすれば、追加・削除・中央値取得がすべて \(O(\log Q)\) でできます。

アルゴリズム

  1. 全クエリを先読みし、出現する値を vals に集める。
  2. vals をソートして重複除去した vals_sorted を作る(座標圧縮)。
    • comp[x] = 圧縮後インデックス(1-indexed)
  3. Fenwick Tree fw を用意し、fw[i] は「圧縮値 i の個数」を表す。
  4. size を現在の要素数として管理。
  5. 各クエリを順に処理:
    • + x:fw.add(comp[x], 1)、size += 1
    • - x:
      1. \(k = (size+1)//2\) を計算(現在の中央値の順位)
      2. mid_idx = fw.kth(k) で \(k\) 番目の圧縮位置を取得
      3. 実値 vals_sorted[mid_idx-1] が x と等しければ答えを +1
      4. その後 fw.add(comp[x], -1)、size -= 1

fw.kth(k) は「累積和が初めて \(k\) 以上になる最小インデックス」を返す関数で、これがそのまま \(k\) 番目の要素に対応します。

計算量

  • 時間計算量: \(O(Q \log Q)\)
    (座標圧縮のソート \(O(Q\log Q)\) + 各クエリの Fenwick 操作 \(O(\log Q)\))
  • 空間計算量: \(O(Q)\)
    (クエリ保存・圧縮配列・BIT)

実装のポイント

  • 中央値は削除前の集合で判定すること(先に判定してから削除)。

  • 問題の中央値定義は「偶数個なら中央2つの小さい方」なので、順位は常に \(k=(size+1)//2\)。

  • Fenwick Tree は 1-indexed で実装すると扱いやすい。

  • kth 実装では BIT の内部配列を使って二分探索的に位置を特定している。これにより \(O(\log N)\) で順位検索が可能。

    ソースコード

import sys


class Fenwick:
    def __init__(self, n):
        self.n = n
        self.bit = [0] * (n + 1)

    def add(self, i, x):
        while i <= self.n:
            self.bit[i] += x
            i += i & -i

    def sum(self, i):
        s = 0
        while i > 0:
            s += self.bit[i]
            i -= i & -i
        return s

    def kth(self, k):
        idx = 0
        bit_mask = 1 << (self.n.bit_length() - 1)
        while bit_mask:
            nxt = idx + bit_mask
            if nxt <= self.n and self.bit[nxt] < k:
                k -= self.bit[nxt]
                idx = nxt
            bit_mask >>= 1
        return idx + 1


def main():
    input = sys.stdin.readline
    Q = int(input())
    queries = []
    vals = []
    for _ in range(Q):
        op, x = input().split()
        x = int(x)
        queries.append((op, x))
        vals.append(x)

    vals_sorted = sorted(set(vals))
    comp = {v: i + 1 for i, v in enumerate(vals_sorted)}

    fw = Fenwick(len(vals_sorted))
    size = 0
    ans = 0

    for op, x in queries:
        idx = comp[x]
        if op == '+':
            fw.add(idx, 1)
            size += 1
        else:
            k = (size + 1) // 2
            mid_idx = fw.kth(k)
            if vals_sorted[mid_idx - 1] == x:
                ans += 1
            fw.add(idx, -1)
            size -= 1

    print(ans)


if __name__ == "__main__":
    main()

この解説は gpt-5.3-codex によって生成されました。

投稿日時:
最終更新: