Atcoder解いてみた AtCoder Beginner Contest C - Drop Blocks

AtCoder Beginner Contest C - Drop Blocks

問題はこちら

考えた方針

素直に問題文に書いてある通りのシミュレーションは行えない
Q回ループは回す必要があるため、O(Q)や、O(Q logN)、O(Q logQ)、O(Q + N)、あたりに収めたい
クエリ1 x, 2 yについて効率的な方法を考える

クエリ2 y

常にソートされているデータ構造を使い、二分探索等を使ってy個以上のブロックが積まれているマスの数を一発で数えられないか
→multisetに(ブロック数, マスの番号)を入れたものを活用できないか

クエリ1 x

x 番目のマスにブロックを 1 個積む

multisetに(ブロック数, マスの番号)を入れたものを活用すればできそう(multisetの中身をeraseし、更新後、りinsertするとか)

すべてのマスに 1 個以上のブロックが積まれているならば、すべてのマスからブロックを 1 個ずつ取り除く

素直に全てのマスから1個取り除くことはできないため、代わりに何回この操作が行われたか記憶し、後でつじつま合わせすればよさそう

このような方針を考え、以下のように実装を進めたが、multisetで二分探索を行ったとしても、multisetのデータ構造上、二分探索でヒットした要素が何番目にあるのかを調べるためには、O(N)かかることを知り、ギブアップ

最初の回答

#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 N;
int Q;

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

	std::cin >> N >> Q;

	// msの初期化
	std::multiset<std::pair<int, int>> ms; // (ブロック数, マス)のms
	for (int n = 0; n < N; n++) {
		std::pair<int, int> pair= std::make_pair(0, n);
		ms.insert(pair); // ブロック数で昇順 ブロック数が同じな場合、マスの番号で昇順
	}
	std::vector<int> blockCount(N);

	int dc = 0; // 全マスのブロック数が-1された回数
	for (int q = 0; q < Q; q++) {
		std::pair<int, int> query;
		std::cin >> query.first >> query.second;

		if (query.first == 1) {
			int x = query.second - 1;
			// 全てのマスが1個以上だったら、全てのマスからブロックを1つのぞく → 結局O(N)かかる?
			// これに関しては、何回、全てに対して-1されたか記録しておきたい

			// x番目のマスのブロック数をインクリメントする前に、x番目のブロックを探し、削除
			std::pair<int, int> pair = std::make_pair(blockCount[x], x);
			auto itr = ms.find(pair);
			if (itr != ms.end()) {
				ms.erase(pair);
			}

			// インクリメント
			pair.first++;
			blockCount[x]++;

			// msに戻す
			ms.insert(pair);

			// 全てのマスが1個以上だったら、全てのマスからブロックを1つのぞく
			// 素直に全マスに-1できないため、全マスのブロック数が-1された回数を使用しているためそれを考慮
			if (ms.begin()->first >= 1 + dc) {
				dc++;
			}
		}
		else if (query.first == 2) {
			// y以上のブロックがあるマスの個数が知りたい → ソート済み配列を使いたい
			// 全マスのブロック数が-1された回数を使用しているためそれを考慮
			int th = query.second + dc;

			// msを二分探索し、
			int idx = ms.lower_bound({ th, INT_MIN }) - ms.begin(); // これができない std::distance使うと、O(N)でNG
			// 0以上は何個のマス、1以上は何個のマスといった形で持っておけば良くないか
			// クエリ与えられたらすぐこたえられるようにする

		}
	}

	return 0;
}

解説の方針

クエリ1 xにおいて、素直に全てのマスから1個取り除くことはできないため、代わりに何回この操作が行われたか記憶し、後でつじつま合わせする方針は合っていた
一方で、クエリ2 yに関しては、このクエリが来た時に、一発で条件を満たす値を取得可能なデータ構造を使用していた
そのデータ構造は、ブロック0個以上のマスの個数, ブロック1個以上のマスの個数...を管理しているもの
自分は、基本的に、データの中から欲しいデータを取得する時、O(logN)くらいかかるだろうと考えてしまっていたが、その他の計算量で取得できないかと疑っていれば、解説に記載されていたデータ構造をひらめくことができたかもしれない

解説をヒントに書いたコード

実装時に注意すべきなのは、masuCountの要素数
仮に、N=1、Q=300,000の時、1種類目のクエリが来るたびに必ずdcがインクリメントされる
そのため、dcは約300,000になる
その上で二種類目のクエリでyに300,000が設定されると、約600,000のインデックスにアクセスすることになる
今後、インデックスに中間データを足したものが使われている場合はしっかり考察すること

#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 N;
int Q;

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

	std::cin >> N >> Q;
	int dc = 0; // 全マスのブロック数が-1された回数

	// 0-indexed
	// masuCountの要素数に注意
	// 仮に、N=1、Q=300,000の時、1種類目のクエリが来るたびに必ずdcがインクリメントされる
	// そのため、dcは約300,000になる
	// その上で二種類目のクエリでyに300,000が設定されると、約600,000のインデックスにアクセスすることになる
	std::vector<int> masuCount(600000); // ブロック0個以上のマスの個数, ブロック1個以上のマスの個数...を管理 このデータ構造だったら常にソートされてる
	std::vector<int> blockCount(300010);

	std::vector<int> ans;

	masuCount[0] = N; // 最初、全てのマスのブロック数が0
	
	for (int q = 0; q < Q; q++) {
		std::pair<int, int> query;
		std::cin >> query.first >> query.second;

		if (query.first == 1) {
			int x = query.second-1; // 番目は0-indexed

			// x番目のマスにブロックを1個積む
			blockCount[x]++;

			// 各ブロック数のマスのカウントを更新
			masuCount[blockCount[x]]++;

			// 全てのマスが1個以上だったら、全てのマスからブロックを1つのぞく
			// 素直に全マスに-1できないため、全マスのブロック数が-1された回数を使用しているためそれを考慮
			if (masuCount[1+ dc] == N) {
				dc++;
			}
		}
		else if (query.first == 2) {
			int y = query.second; // y>=1 ブロックの個数はすでに0-indexed

			// 素直に全マスに-1できないため、全マスのブロック数が-1された回数を使用しているためそれを考慮
			ans.push_back(masuCount[y+dc]);
		}
	}

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

AIのフィードバック

• std::vector masuCount(Q + 300010 + 5);とした方がコードの意図が伝わりやすい(+5は少し余裕を持たせるため)
• 「ある値以上の個数」という時点で、ソート集合,累積個数配列あたりの選択肢を持った上で、ブロックの個数の単調性(ただ積まれるだけ)に着目できればよかったかもしれない