Editorial - AtCoder Weekday Contest 0055 Beta

Official

D - 中央値チェック / Median Check Editorial by admin

GPT 5.4 High

概要

多重集合に対する追加・削除を順に処理し、削除される値が「削除直前の中央値」と一致していた回数を数える問題です。
各時点での中央値を高速に求めるために、座標圧縮 + Fenwick Tree（BIT） を使って管理します。

考察

この問題では、削除クエリ - X のたびに

削除直前の多重集合の中央値を求める
それが \(X\) と一致するか判定する

必要があります。

中央値は「何番目の値」か

要素数を \(n\) とすると、問題文で定義される中央値は

昇順に並べたときの \(\lceil n/2 \rceil\) 番目
これは整数除算で書くと \(((n+1)//2)\) 番目

です。

例えば：

\(n=1\) のとき → 1 番目
\(n=3\) のとき → 2 番目
\(n=4\) のとき → 2 番目（真ん中2つのうち小さい方）

つまり、「現在の多重集合の \(k=((n+1)//2)\) 番目に小さい値」が中央値です。

素朴な方法が難しい理由

多重集合を毎回ソートしたり、配列に挿入・削除したりすると重くなります。

例えば、昇順配列で管理すると：

挿入位置を二分探索で見つけても、実際の挿入・削除は \(O(n)\)
これを \(Q\) 回繰り返すと最悪 \(O(Q^2)\)

\(Q \le 2 \times 10^5\) なので、これは間に合いません。

必要な操作は何か

この問題で欲しいのは、各時点で次の 2 つです。

値 \(X\) を 1 個追加 / 削除する
「\(k\) 番目に小さい値」を求める

これは、各値が何個あるか を管理できれば解けそうです。

値が大きいので座標圧縮する

\(X_i\) は最大 \(10^9\) ですが、実際に登場する値は高々 \(Q\) 個です。
そこで、登場する値を昇順に並べて

小さい順に 1, 2, 3, … と番号を振る

ことで 座標圧縮 します。

例えば値が \(\{100, 500, 1000\}\) しか出ないなら、

\(100 \to 1\)
\(500 \to 2\)
\(1000 \to 3\)

と変換できます。

この圧縮後の添字に対して個数を管理すればよいです。

Fenwick Tree で個数管理

Fenwick Tree（BIT）で、圧縮後の各値の出現回数を持ちます。

すると、

add(idx, 1) で値を 1 個追加
add(idx, -1) で値を 1 個削除
sum(i) で「圧縮値が \(i\) 以下の要素数」が求まる

ようになります。

ここで、中央値は「\(k\) 番目に小さい値」なので、

累積個数が初めて \(k\) 以上になる位置

を探せばよいです。

これを Fenwick Tree 上で行うのが lower_bound(k) です。

どうしてこれで中央値が分かるのか

たとえば多重集合が

\(\{2, 2, 5, 8, 8\}\)

だとします。要素数は 5 なので、中央値は \(((5+1)//2)=3\) 番目の値です。
昇順で 3 番目は 5 です。

個数で見ると：

2 が 2 個
5 が 1 個
8 が 2 個

累積個数は

2 以下：2 個
5 以下：3 個
8 以下：5 個

このとき、累積個数が初めて 3 以上になるのは 5 の位置なので、中央値は 5 と分かります。

削除クエリでやること

- X のときは、削除する前に 中央値を求めます。

現在の要素数を total とすると、

中央値の順位は \(k = (total + 1) // 2\)
Fenwick Tree で \(k\) 番目の値の圧縮添字 med_idx を求める
削除対象の圧縮添字 idx と比較する
一致すれば「中央値ヒット」

その後で実際に削除します。

この「判定してから削除」の順番が重要です。

アルゴリズム

すべてのクエリを先に読み込み、登場する値を集める
値を座標圧縮する
Fenwick Tree を用意し、各圧縮値の個数を管理する
total を現在の要素数として持つ
各クエリを順に処理する
- + X
  - 対応する圧縮位置 idx に 1 を加える
  - total += 1
- - X
  - 中央値の順位 \(k = (total + 1) // 2\) を求める
  - Fenwick Tree の lower_bound(k) で中央値の圧縮位置 med_idx を求める
  - med_idx == idx なら答えを 1 増やす
  - その後 idx の個数を 1 減らす
  - total -= 1
最後に答えを出力する

`lower_bound(k)` とは

Fenwick Tree の lower_bound(k) は、

累積和が初めて \(k\) 以上になる最小の添字

を返します。

これはつまり、\(k\) 番目に小さい値の位置 です。

この問題ではそれがそのまま中央値の位置になります。

計算量

時間計算量: \(O(Q \log Q)\)
空間計算量: \(O(Q)\)

補足

座標圧縮のためのソートが \(O(Q \log Q)\)
各クエリの追加・削除・中央値取得がそれぞれ \(O(\log Q)\)

実装のポイント

中央値は削除前に判定する
先に削除してしまうと、比較するべき中央値が変わってしまいます。
偶数個のときも \(k=(n+1)//2\) でよい
問題文では「真ん中2つのうち小さい方」なので、例えば \(n=4\) なら 2 番目の値です。
値そのものではなく圧縮後の添字で比較してよい
座標圧縮は大小関係を保つので、中央値の値の比較は圧縮添字同士で問題ありません。
Fenwick Tree は 1-indexed で実装する
このコードでも圧縮値を i + 1 にしており、Fenwick Tree と相性がよくなっています。

ソースコード

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 lower_bound(self, w):
        if w <= 0:
            return 1
        x = 0
        k = 1 << (self.n.bit_length() - 1)
        while k:
            nx = x + k
            if nx <= self.n and self.bit[nx] < w:
                w -= self.bit[nx]
                x = nx
            k >>= 1
        return x + 1

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

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

    fw = Fenwick(len(vals))
    total = 0
    ans = 0

    for op, x in queries:
        idx = comp[x]
        if op == '+':
            fw.add(idx, 1)
            total += 1
        else:
            k = (total + 1) // 2
            med_idx = fw.lower_bound(k)
            if med_idx == idx:
                ans += 1
            fw.add(idx, -1)
            total -= 1

    print(ans)

if __name__ == "__main__":
    main()

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