Atcoder解いてみた AtCoder Beginner Contest D - Chalkboard Median

AtCoder Beginner Contest D - Chalkboard Median

問題はこちら

回答

素直にやるとすると、クエリの度に数値をソートして中央値を求める
その場合、計算量がO(Q Q logQ)となるため、アウト
少なくとも、毎回ソートを行わないようにし、計算量をO(Q logQ)あたりにしたい
multisetや、priority_queueをうまく使えないかと思ったが、思いつかず、、、
調べてみると、priority_queueを2つ使うというキーワードが出てきたため、それをヒントに実装できた
計算量はO(Q logQ)
multiset使っても解くことができる

#include <iostream>
#include <string>
#include <map>
#include <unordered_map>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <vector>
#include <queue>
#include <stack>
#include <limits.h>
#include <bitset>
#include <list>
#include <set>
#include <numeric>
#include <tuple>

int X;
int Q;

std::priority_queue<int> pqWithMed;
std::priority_queue<int, std::vector<int>, std::greater<int>> pqWithOutMed; // 中央値以上を格納するため、topが中央値となるようにする


int main()
{
	std::cin.tie(0);
	std::ios::sync_with_stdio(false);

	std::cin >> X;
	std::cin >> Q;

	pqWithMed.push(X);

	std::vector<int> ans;
	for (int i = 0; i < Q; i++) {
		std::vector<int> AB(2);
		std::cin >> AB[0] >> AB[1];

		for (int j = 0; j < AB.size(); j++) {
			// pqへの格納
			// 現在の中央値Xが入力された値以上の時、pqWithMedに入力された値をpush
			if (X >= AB[j]) {
				pqWithMed.push(AB[j]);
			}
			else {
				pqWithOutMed.push(AB[j]);
			}
		}

		// pqの調整
		while (pqWithMed.size() < pqWithOutMed.size() + 1) {
			pqWithMed.push(pqWithOutMed.top());
			pqWithOutMed.pop();
		}

		while (pqWithMed.size() > pqWithOutMed.size() + 1) {
			pqWithOutMed.push(pqWithMed.top());
			pqWithMed.pop();
		}

		X = pqWithMed.top(); // 中央値の更新
		ans.push_back(X);
	}

	for (int i = 0; i < ans.size(); i++) {
		std::cout << ans[i] << std::endl;
	}

	return 0;
}