AtCoder Beginner Contest 260, D : Draw Your Cards

問題文

　https://atcoder.jp/contests/abc260/tasks/abc260_d

問題概要

　$1$ 以上 $N$ 以下の整数が重複なく書かれたカードが $N$ 枚あり，山札として積まれている．上から $i$ 番目のカードに書かれているのは $P_i$ である．
　ここで，次の操作を $N$ 回繰り返す．

1. 山札の一番上のカードを引く．それに書かれていた整数を $X$ とする．
2. 場に見えている（後述）表向きのカード（後述）であって，書かれているのが $X$ 以上であり，その内で書かれている整数が最小のカードに，引いたカードを表向きで重ねて置く（重ねられたカードは見えなくなる）．重ねるべきカードが存在しない場合は，引いたカードを単体で（表向きで）場に置く．
3. 表向きのカードが $K$ 枚以上重ねられている箇所が場にあれば，それを除外する．

　各カードについて，それが除外されたのが何回目の操作時であるか報告せよ．除外されない場合は，その旨を報告せよ．

制約

• $1 \leq K \leq N \leq 2 \times 10^5$
• $P$ はサイズ $N$ の順列

解法

　場に表向きで積まれているカードの束それぞれを「山」と呼ぶことにします（場には山の集合があることになります）．
　色々と難しいので，なんとかシミュレーションできないか考えてみましょう．つまり，山の集合を管理して，それを適切に（矛盾無く，TLE しない程度に高速に）更新することを考えます．
　使用するプログラミング言語に依存する話になってしまいますが，各山の枚数は動的に変動し，山の数自体も動的に変動するので，二次元にした std::vector を使うことにします（C++ の話．他の言語を使う場合は同等のデータ構造（動的配列的なやつ）で読み替えてください）．この二次元 vector の名前を $D$ とします*1．また，vector の末尾を参照する関数を $\mathit{ back }( * )$ とします．このとき，各操作で $X$ が定まったときにやりたいことは，

1. $X \leq \mathit{ back }( D_i )$ な $i$ のうち，$\mathit{ back }( D_i )$ が最小となる $i$ を見つける．存在しない場合は $D$ の末尾に新たな（空の）配列を追加し，それを指す添字を $i$ とする．
2. $D_i$ の末尾に $X$ を追加する
3. $K \leq | D_i |$ となったなら，$D_i$ の要素それぞれについて，今が何回目の操作であったかを記録しておく．

です．時間がかかりそうなのは 1. と 3. ですが，3. については，カードに書かれた各整数それぞれについて高々一回しか操作の対象にならないため，全体で $O( N )$ 時間ということになって問題ありません．1. についてはもうちょっと考える必要があって，山の数が $O( N )$ 個まで増えるので，毎回 $O( N )$ 時間かけて探したりすると全体 $O( N^2 )$ 時間で TLE します．全体で $\Theta( N )$ 回の操作は確定なので，各操作をシミュレートする時間を $o( N )$ 時間……というか，やや天啓かつ典型ですが $O( \log N )$ 時間にしたい気持ちになります．
　そこで，操作対象となる山を指す添字 $i$ をより高速に見つけるため，別のデータ構造を用意します．$X$ と $\mathit{ back }( D_i )$ の関係についておさらいすると，$X$ 以上で最小の $\mathit{ back }( D_i )$ が欲しいわけですから，$\mathit{ back }( D_i )$ が昇順に並んだ列があれば，その上で二分探索をすることで該当する $\mathit{ back }( D_i )$ を見つけられます．もっと言えば，$\mathit{ back }( D_i )$ と $i$ を順序対にしたものを（順序対の辞書順比較という意味での）昇順に並べておけば，その列上で $( X, -\infty )$ 以上の順序対となる最初の値が，欲しかった情報です．この順序対の列に対して列への新たな要素の挿入・削除と列上の二分探索が $O( \log N )$ 時間となるデータ構造が欲しいわけですが，これには std::set が利用できます．この set を $S$ とします．
　これで，さっきの「やりたかったこと」を実現する手順を書けるようになりました．具体的には，

1. カードを一枚引き，それに書かれた数を $X$ とする．
2. $S$ 上の二分探索により，$( X, -\infty )$ となる最小の順序対 $( x, i )$ を求める（この第二要素が操作対象になる山の添字）．存在しない場合は，代わりに $( x, | D |)$ とした上で $D$ の末尾に空の列を追加する（添字は 0-indexed）．存在したなら，今見つけた要素を $S$ から削除する．
3. $D_i$ の末尾に $X$ を追加する．ここで，$K \leq | D_i |$ であれば，$D_i$ の各要素について答えを集計する．そうでない場合，$S$ に $( X, i )$ を追加する．

のようにします．
これで $\Theta( N )$ 回の操作をそれぞれ $O( \log N )$ で処理できるようになり，全体で $O( N \log N )$ 時間となって間に合います．

コード

constexpr auto INF = LIM<>::max() / 2;

int main()
{
	cin.tie( nullptr );
	ios::sync_with_stdio( false );
	cout << setprecision( 12 ) << fixed;

	IN( int, N, K );
	VI P( N );
	cin >> P;

	VVI decks;
	decks.reserve( N );
	set< PII > indices;
	VI res( N, -1 );

	REP( t, N )
	{
		const int X = P[t];

		const auto it = indices.lower_bound( MP( X, -INF ) );
		int idx = -1; 
		if ( it == end( indices ) )
		{
			idx = SZ( decks );
			decks.EB();
		}
		else
		{
			idx = it->snd;
			indices.erase( it );
		}
		decks[ idx ].PB( X );

		if ( K <= SZ( decks[ idx ] ) )
		{
			FOR( p, decks[ idx ] )
			{
				res[ p - 1 ] = t + 1;
			}
		}
		else
		{
			indices.EM( X, idx );
		}
	}

	copy( ALL( res ), OSI< int >( cout, "\n" ) );
	cout << flush;

	return 0;
}

*1:なんとなく Decks からとったつもり