AtCoder Beginner Contest 392, D : Doubles

問題文

　https://atcoder.jp/contests/abc392/tasks/abc392_d

問題概要

　$n$ 個のダイスがある．$i$ 番目のサイコロは $K_i$ 面ダイスであり，その各面の出目は $A_{ i, 1 }, A_{ i, 2 }, \dots, A_{ i, K_i }$ である．ダイス $i$ を振ったとき，各出目が出る確率は等しく $\frac 1 { K_i }$ である．
　$n$ 個のダイスから $2$ 個を選んで同時に振ることを考える．$2$ つの出目が等しくなる確率はダイスの選び方によって異なるが，確率を最大化した場合のその確率はいくらか？

制約

• $2 \leq n \leq 100$
• $1 \leq K_i$
• $\sum K \leq 10^5$
• $1 \leq A_{ i, j } \leq 10^5$

解法

　ダイスの組合せの内で考慮しないものが存在してしまうと最大値を正しく計算できないように思えるので，ダイスの組合せはすべて試すことにします．選ぶダイス $i, j$ ($i \neq j$) を決めたときに出目が揃う確率をすべての $i, j$ について求めることができれば，それらの最大値を出力することで問題を解くことができます．
　振るダイス $i, j$ を固定したときに求めたい確率は，素朴には
\begin{equation}
\frac{ | \{ ( x, y ) \mid
x \in \{ 1, 2, \dots, K_i \},
y \in \{ 1, 2, \dots, K_j \},
, A_{ i, x } = A_{ j, y } \} | }
{ K_i K_j }
\end{equation}
ですが，これをそのまま計算するために $x, y$ の組合せをすべて試すと TLE するので工夫します．ありがちな戦略として，$x, y$ の片方だけを固定してみます．つまり，ダイス $i$ の出目 $A_{ i, x }$ を固定したときに $A_{ i, x } = A_{ j, y }$ となるような $y$ の個数を愚直にやるより高速に求めたいです．やや天啓ですが，このとき $A_j$ が昇順ソートされていれば，$A_j$ の添字 $l, r$ を
\begin{align*}
l &= \text{$A_{ i, x } \leq A_{ j, y }$ となるような最小の $y$} \\
r &= \text{$A_{ i, x } < A_{ j, y }$ となるような最小の $y$}
\end{align*}
のようにとることで，求めたかった個数を $r - l$ として計算できます．また，$l, r$ は二分探索により高速に求めることができます．更に，$A_i$ も昇順にソートされていれば，$x$ を昇順に動かすに伴って $l, r$ は非減少に動くので，$x$ が変化した後で $l, r$ が条件を満たす位置になるまでインクリメントすることで，$l, r$ が変化する回数を $\Theta( K_j )$ 回にすることができます．まとめると，各 $A_i$ をあらかじめソートしておけば，以下のようなコードによってダイス $i, j$ を振って出目が揃う確率を $\Theta( K_i + K_j )$ 時間で求めることができます（このコード片の外側で $i, j$ の二重ループが回ります）：

double p = 0;
for ( int k = 0, l = 0, r = 0; k < SZ( A[i] ); ++k )
{
	for ( ; l < SZ( A[j] ) && A[j][l] < A[i][k]; ++l );
	for ( ; r < SZ( A[j] ) && A[j][r] <= A[i][k]; ++r );
	p += 1. / SZ( A[i] ) * ( r - l ) / SZ( A[j] );
}

　では，計算量はどうなっているでしょうか．素朴に考えると，$1, 2, \dots, n$ の範囲をそれぞれ動く $i, j$ の二重ループの中で $K_i, K_j$ に関するループが回っていて，上界を求めたい気持ちで最大値をとると $O( n^2 \sum K )$ 時間かかるように見えます．（言語によって？）ギリギリ通りそうなところではありますが，実は解析を工夫するだけでよりタイトな上界を得られます．
　上で示したコード片において，$A$ の各要素がループ変数 $k$ によって参照される回数を考えると，各 $i$ について $\Theta( n )$ 回ずつ（$i$ に対応してとれる $j$ の個数に対応）です．よって，$A$ 全体では $\Theta( n \sum K )$ 回になります．同様にループ変数 $l, r$ によって参照される回数を考えると，各 $j$ に対して $\Theta( n )$ 回ずつ（$j$ に対応してとれる $i$ の個数に対応）なので，こちらも $A$ 全体で $\Theta( n \sum K )$ 回になります．従って，全体では前処理のソートと合わせて $O( \sum K \log \sum K + n \sum K )$ 時間となります．

コード

int main()
{
	IN( int, N );
	VVI A;
	generate_n( BI( A ), N, []{
			IN( int, K );
			VI row( K );
			cin >> row;
			sort( ALL( row ) );
			return row; } );

	double res = 0;
	REP( i, N )
	{
		REP( j, i )
		{
			double p = 0;
			for ( int k = 0, l = 0, r = 0; k < SZ( A[i] ); ++k )
			{
				for ( ; l < SZ( A[j] ) && A[j][l] < A[i][k]; ++l );
				for ( ; r < SZ( A[j] ) && A[j][r] <= A[i][k]; ++r );
				p += 1. / SZ( A[i] ) * ( r - l ) / SZ( A[j] );
			}
			chmax( res, p );
		}
	}

	cout << res << endl;

	return 0;
}