SRM 370 250pt: DrawingMarbles

問題概要

箱にボールがいくつか入っている。色はN(<50)色で、各色はそれぞれA[i](<50)個ずつ入っている。この中からK個取り出すとき、全ての色が異なる確率を求める問題。

考えたこと

  • 高校数学みたいな問題だ。やるだけに見える。
  • 一応オーバーフローとか誤差とかちょっと怖いので対数で計算しよう。
  • 通った。
#include <cmath>
#include <vector>
#include <numeric>

using namespace std;

struct DrawingMarbles {

	double sameColor(vector <int> colors, int n) {
		const int N = colors.size();
		const int total = accumulate(colors.begin(), colors.end(), 0);

		double ret = 0.0;
		for(int i=0; i<N; i++)if(colors[i] >= n){

			//誤差があるので対数で計算
			double lg = 0.0;

			//分母
			for(int j=0; j<n; j++){
				lg -= log(total - j);
			}

			//分子
			for(int j=0; j<n; j++){
				lg += log(colors[i] - j);
			}

			ret += exp(lg);
		}

		return ret;
	}

};

他の人の解法

オーバーフロー対策で、logとか使わなくてもよかった。分子だけ、分母だけを先に計算するからまずいのであって、今回の場合は同時にかけたり割ったりすれば大丈夫だった。こんな風に。

		double ret = 0.0;
		for(int i=0; i<N; i++)if(colors[i] >= n){

			double prob = 1.0;
			for(int j=0; j<n; j++){
				prob *= (double)(colors[i] - j)/(total - j);
			}

			ret += prob;
		}

それ以外にも、途中で値が大きくなりすぎたらその時点で割り算するという方法もある。以前POJでやったことがあるような。