AtCoder Beginner Contest 380, E : 1D Bucket Tool

問題文

　https://atcoder.jp/contests/abc380/tasks/abc380_e

問題概要

　$1$ から $n$ の整数で番号付けられたセルがあり，セル $i$ とセル $i + 1$ ($1 \leq i < n$) は隣接している．各セル $i$ は色 $i$ で塗られている．
　以下の 2 種類のクエリが $q$ 個与えられるので，順に処理せよ：

• クエリ $1$ : 整数 $x, c$ が与えられる．セル $x$ から始めて同色で塗られたセルへの移動を繰り返して到達できるセルすべてを色 $c$ で塗り替える．
• クエリ $2$ : 整数 $c$ が与えられる．色 $c$ で塗られたセルの個数を答える．

制約

• $1 \leq n \leq 5 \times 10^5$
• $1 \leq q \leq 2 \times 10^5$

解法

　クエリ $1$ で色が変わるセルの個数は $\Theta( n )$ 個になり得ます．悪意があるケースを愚直に処理すると $\Omega( q n )$ 時間がかかって TLE するので，工夫する必要があります．
　クエリ $1$ で同時に色が変わるセルは，（整数の）区間 $[ 1, n ]$ に内包される区間になっています．それぞれのセルをバラバラに管理する代わりに同色で塗られた（極大な）区間を管理することで，色の変化を一括で処理することができそうです．色の情報を同時に管理するために，区間を triplet $( l, r, c )$ で表します．意味は

• 区間 $[ l, r )$ は色 $c$ で塗られている．

です．
　上記 triplet を何らかのデータ構造 $\mathcal S$ に入れて管理します．$\mathcal S$ の具体的な実装として何を選ぶ（あるいは新たに実装する）べきかを判断するために，$\mathcal S$ に対してどのような操作が必要となるかを考えます．クエリ $1$ を処理するときにやるべきことは，

1. $x$ を含む（唯一の）区間を $\mathcal S$ から見つけ，$T$ とする．
2. $T$ を $\mathcal S$ から削除する．
3. $T$ の直前・直後の区間を表す triplet をそれぞれ $( l_l, r_l, c_l ), ( l_r, r_r, c_r )$ とし，$c$ に一致する $c_l, c_r$ があればそれを含む triplet を $\mathcal S$ から削除する．
4. 削除された区間たちをちょうど被覆する区間を $[ l, r )$ とし，$( l, r, c )$ を $\mathcal S$ に追加する．

となります．1. の実装に効率的な探索方法を採用する必要がありますが，triplet は自然な辞書式順序で順序付けることができるため，二分探索を用いることができます．加えて要素の追加・削除も効率的に行う必要がありますが，二分探索と同様のアイデアに基づく探索と追加・削除ができるするデータ構造としては，平衡二分探索木が有名です．C++（の典型的な実装）では std::set が利用できます．$x$ を含む区間は std::lower_bound で $( x + 1, 0, 0 )$ などを探し，その直前の triplet を見ると欲しかったものになっています．std::prev, std::next を使うことで，参照したい triplet を指すイテレータはすべて得られるので，上記の手続きを実装できます．
　加えて，クエリ $2$ に答えるために，色ごとにその出現数を管理する配列を用意します．triplet の追加・削除を行うたびにその幅分増減させることで，$O( 1 )$ 時間でクエリに答えることができます．
　上記手続きはクエリあたり $O( \log n )$ 時間で実行できるため，全体では $O( q \log n )$ 時間となります．

実装について

　実装するにあたっては，区間 $[ 0, 1 ), [ n, n + 1 )$ を存在しない色で塗ったものとして $\mathcal S$ に入れて番兵にすることで細かい分岐処理を消去できます．特に，std::prev でイテレータがはみ出してしまう*1ことを防ぐのが簡単になります．

コード

int main()
{
	IN( int, N, Q );
	
	using Interval = tuple< int, int, int >;
	// ( l, r, c ) denotes [ l, r ) that colored c
	set< Interval > intervals;
	REP( i, -1, N + 1 )
	{
		intervals.EM( i, i + 1, i + 1 );
	}
	VI counts( N + 2, 1 );

	REP( Q )
	{
		IN( int, TYPE );
		if ( TYPE == 1 )
		{
			IN( int, X, C );
			--X;

			const auto it = --intervals.lower_bound( Interval( X + 1, 0, 0 ) );
			const auto [ l, r, c ] = *it;
			const auto itp = prev( it ), itn = next( it );
			int nl = l, nr = r;

			static const auto remove = [&]( const auto it, const bool force = false )
			{
				const auto [ ll, rr, cc ] = *it;
				if ( !force && cc != C )
				{
					return;
				}
				chmin( nl, ll );
				chmax( nr, rr );
				intervals.erase( it );
				counts[ cc ] -= rr - ll;
				return;
			};

			remove( it, true );
			remove( itp );
			remove( itn );
			intervals.EM( nl, nr, C );
			counts[C] += nr - nl;
		}
		else
		{
			IN( int, C );

			cout << counts[C] << '\n';
		}
	}

	cout << flush;

	return 0;
}

*1:未定義動作になる気がする