SRM 462 450pt: CandyBox

TopCoder

問題概要

N(<50)個の箱がある。それぞれの箱にはそれぞれ同じ種類のアメがC(<100)個ずつ入っている。あなたのアメの好みはscore[i](<100)で表される。さて、無作為に2つのアメを選んで交換するという作業をS(<10^4)回繰り返す。その後各箱に対して、1個無作為に選んだとき得られるスコアの期待値をそれぞれ求める問題。

考えたこと

  • (当時は問題の意味くらいは分かったけど解き方は全然思い浮かばなかった)
  • Cとかscoreが小さいしDP？
  • 1個無作為に選んだときの期待値じゃなくて、S回スワップした後の箱の中のアメのスコアの総和を計算するようにしよう。
  • ある箱のスコアの和が分かったら、それ以外の箱の和も引き算で求まる。これは使えそう。
  • というか、各箱についてまったく独立に計算できる。
  • 状態数をターン数にして、箱の中のアメのスコアの和の期待値でDPいけそう。
  • スワップするときある特定の列が選ばれる確率とかは前もって計算しておける。
  • 書く。サンプル通らないし、総和が減ってておかしい。
  • 期待値の計算式に間違い発見。訂正。まだ合わない。
  • その列がスワップに含まれる確率を訂正。サンプル通った。
  • 提出。無事通った。以前は手も足もでなかった問題が通ると無性に嬉しい。
  • 続きの部分に書いてるけど、この問題は計算量が落とせて色々楽しい。

practiceで通ったコード

計算量O(N*S)。

#include <vector>
#include <numeric>
using namespace std;


struct CandyBox {

	vector <double> expectedScore(int C, vector <int> score, int S) {
		const int N = score.size();

		//コーナーケース
		if(N==1){
			return vector<double>(score.begin(), score.end());
		}

		int acc_score = accumulate(score.begin(), score.end(), 0) * C;

		//ある特定の列で入れ替わりが生じる確率
		double prob_change = (double)2.0 * C * (N-1) * C / (N*C) / (N*C-1);

		vector<double> ret(N, 0.0);
		for(int i=0; i<N; i++){
			//expectは持ってる量の期待値
			double& expect = ret[i];
			expect = score[i] * C;

			//スワップによって期待値が変動する
			for(int j=0; j<S; j++){
				expect = prob_change * ( (acc_score - expect)/(C*(N-1)) - (expect / C) + expect)
						+ (1.0 - prob_change) * expect;
			}
			expect /= C;
		}

		return ret;
	}

};

計算量を落とす

  • スワップによって元のアメがいくつ残るか、を計算しておけばもっと話は単純になる。
  • その期待値は線型性を使って簡単に計算できる。
  • これだと計算量がO(S + N)に落ちる。
#include <vector>
#include <numeric>
using namespace std;

struct CandyBox {

	vector <double> expectedScore(int C, vector <int> score, int S) {
		const int N = score.size();

		//コーナーケース(0割逃れ)
		if(N == 1){
			return vector<double>(score.begin(), score.end());
		}

		//総和を前処理で持っておく
		int acc_score = accumulate(score.begin(), score.end(), 0);

		//ある1個がそのままの箱に止まっている確率
		double stay_prob = 1.0;

		double away_prob = (double)((N-1)*C) / ((double)(N*C)*(N*C-1)/2.0) , //違う箱へ行ってしまう確率
				back_prob = (double)C / ((double)(N*C)*(N*C-1)/2.0); //違う箱から戻ってくる確率
		for(int i=0; i<S; i++){
			stay_prob = stay_prob * (1.0 - away_prob) + (1.0 - stay_prob) * back_prob;
		}
		//もとの箱に残っている個数の期待値
		double expect_left = stay_prob * C; //期待値の線型性


		//答えの計算
		vector<double> ret(N, 0.0);
		for(int i=0; i<N; i++){
			ret[i] = (expect_left * score[i] + (C - expect_left) * (acc_score - score[i]) / (N-1) ) / C;
		}

		return ret;
	}

};

さらに計算量を落とす

  • 上のコード中で、ある特定の1個がそのままの箱にとどまる確率stay_probの計算式を見れば分かるように、これは漸化式になっている。
    • a[i+1] = a[i] * x + (1 - a[i]) * y = (x - y) * a[i] + y
  • これは行列累乗でO(log S)で求めることができる。全体の計算量はO(log S + N)。
  • さらに、この隣接二項間漸化式は手で解くことができる(高校とかで習うアレ)。一般項に行列の巾が出てくるとはいえ、std::pow(double, double)の計算量をO(1)だと思うと全体の計算量はO(N)となりSが無視できる様になる。