Official

D - Chalkboard Median Editorial by sounansya


黒板に書かれた整数の多重集合を \(S\) とします。\(S\) に対する要素の追加クエリが \(Q\) 回くるので、それぞれに対して \(S\) の中央値を求めれば良いです。

ここで、\(S\) を以下の条件を満たすように \(2\) つの多重集合 \(S_1,S_2\) に分けることを考えます。

  • \(|S_1|-|S_2|\)\(0\) または \(1\) である。
  • \(x \in S_1,\ y\in S_2\) に対し \(x \geq y\) が成り立つ。

直感的に言うと、\(S\) を大きい半分 \(S_1\) と小さい半分 \(S_2\) に分割するということです(ただし \(|S|\) が奇数である場合は \(S_1\) の方が要素が \(1\) つだけ多くなるようにする)。

このように分割すると、\(S\) の中央値は \(S_1\) の最小値と一致することが分かります。

さらに、\(S\)\(x\) を追加するようなクエリは \(S_1,S_2\) に対して以下のように作用させることで対応できます:

  • \(x\)\(\min S_1\) 以上である場合、\(S_1\)\(x\) を追加する。そうでない場合、\(S_2\)\(x\) を追加する。(これにより \(x \in S_1,\ y\in S_2\) に対し \(x \geq y\) は保たれたままであるが、\(|S_1|-|S_2| \in \lbrace 0,1\rbrace\) は成り立たない場合がある。)
  • \(|S_1| > |S_2|+1\) である場合、\(\min S_1\)\(S_1\) から \(1\) つ削除し、代わりに \(S_2\) に追加する。
  • \(|S_2| > |S_1|+1\) である場合、\(\max S_2\)\(S_2\) から \(1\) つ削除し、代わりに \(S_1\) に追加する。

これらの操作は \(S_1,S_2\) をそれぞれ優先度付きキューで表現することで簡単に実装することができます。

以上を適切に実装することでこの問題に正答することができます。計算量は \(O(Q \log Q)\) です。

実装例(C++)

#include <bits/stdc++.h>
using namespace std;

int main() {
	ios::sync_with_stdio(false);
	cin.tie(nullptr);
	int x;
	cin >> x;
	priority_queue<int> L;
	priority_queue<int, vector<int>, greater<int>> R;
	L.push(x);
	auto add = [&](int val) {
		if (val <= L.top()) {
			L.push(val);
		} else {
			R.push(val);
		}
		if (L.size() < R.size() + 1) {
			L.push(R.top());
			R.pop();
		}
		if (L.size() > R.size() + 1) {
			R.push(L.top());
			L.pop();
		}
	};
	int Q;
	cin >> Q;
	while (Q--) {
		int a, b;
		cin >> a >> b;
		add(a), add(b);
		cout << L.top() << '\n';
	}
	return 0;
}

実装例(Python3)

import sys

input = sys.stdin.readline
from heapq import heappush, heappop

L = [-int(input())]
R = []


def add(x):
    if x <= -L[0]:
        heappush(L, -x)
    else:
        heappush(R, x)

    if len(L) < len(R) + 1:
        heappush(L, -heappop(R))
    if len(L) > len(R) + 1:
        heappush(R, -heappop(L))


for _ in range(int(input())):
    a, b = map(int, input().split())
    add(a), add(b)
    print(-L[0])

posted:
last update: