TopCoder SRM 577 Div1 Easy EllysRoomAssignmentsDiv1

問題

n人が参加するプログラミングコンテストで、部屋の割り当てがある。
部屋の数はR = n / 20（小数点以下切り上げ）で、
次のように割り当てられる。

n人をレート順にソートする。
大きい順にR人ずつ取り、各人をR個の部屋にランダムに割り当てる。
（二人が一つの部屋に入ることはない）
これを全員が割り当てられるまで繰り返す。

最後はちょうどR人を取ることができない場合があるが、
同じように一人高々一つの部屋をランダムで割り当てる。

n人の中の一人はellyであり、ellyのレートが与えられる。
ellyの部屋の平均レートの期待値を求めよ。

制約条件

n≦500くらい
どの二人のレートも異なる。

方針

nが20で割り切れる場合は、簡単。
ellyと一緒にR人に選ばれる人は同じ部屋になれず、
それ以外の人は確率1/Rで同じ部屋になる。
なので合計レートの期待値はそれらを足せばよく、人数はn / R人なので、
それで和を割れば、平均レートの期待値が求まる。

ellyが最後の余りになる場合とそうでない場合で場合わけ。
最後の余りになる場合、ellyの部屋の人数はn / R + 1

そうでない場合は、(余り) / Rの確率でn / R + 1人部屋になり、
1 - （余り）/ Rの確率でn / R人部屋になる。
それぞれ、同部屋の人のレートの合計の期待値は最初の場合と同じような計算で求められるので、求めて期待値を計算すればよい。

ソースコード

class EllysRoomAssignmentsDiv1 {
	public:
	double getAverage(vector <string> ratings) {
		vi rs;
		{
			string s;
			int x;
			rep(i, ratings.size()) s += ratings[i];
			stringstream ss(s);
			while(ss >> x) rs.pb(-x);
		}
		int elly = rs[0];
		int n = rs.size();
		int R = (n + 19) / 20;
		sort(all(rs));
		
		int idx = lower_bound(all(rs), elly) - rs.begin();
		double res = -elly;
		bool found = 0;
		
		for(int i = 0; i < n / R * R; i += R){
			
			bool isE = 0;
			rep(j, R) if(rs[i + j] == elly) isE = found = 1;
			if(isE) continue;
			rep(j, R) res -= rs[i + j] * 1.0 / R;
		}
		if(n % R == 0){
			return res / (n / R);
		}
		
		if(found){ //エリーが余りに入ってない場合
			double sum1 = res, sum2 = res;
//n/R人部屋とn/R+1人部屋の場合がありうる
			double p = 1.0 * (n % R) / R;
			
			for(int i = n / R * R; i < n; i++) sum2 -= rs[i] * 1.0 / (n % R);
			return sum1 * (1 - p) / (n / R) + sum2 * p / (n / R + 1);
		}
		else{ //エリーが余りに入っている場合
			return res / (n / R + 1);
		}
		return -1;
	}
};