CodeIQ:【CodeIQ数学ゼミ】参加応募用問題

私自身が表題の問題を解いた時のプログラムについて解説します。
問題の詳細は「【CodeIQ数学ゼミ】参加応募用問題」(CodeIQ)を参照してください。

問題の概要

問題を引用します。
箱Aには赤玉が1個と白玉がn-1個、箱Bには赤玉が0個と白玉がm個入っている。
箱A, Bからそれぞれ1個の玉を同時に取り出し、Aから出た玉はBに、Bから出た玉はAに移し替える試行を繰り返す。
この試行をN回繰り返した時、赤玉が箱Aに入っている確率は?

n,m,Nの各データが標準入力でスペース区切りで与えられた際、正解(結果)を標準出力に出力するプログラムを作成してください。

n,m,Nはいずれも10000以下の正の整数で、n,mは1.0のように与えられます。値の精度はn, mはfloatを使用してください。

n m N
のように標準入力が与えられるものとします。

入力例:
1.0 1.0 10

出力例:
1.0

入力例:
1.0 100.0 1000

出力例:
0.00990099009901

私のプログラム

Pythonで解答しています。

	#!/usr/bin/python
	# _*_ coding: utf-8 _*_

	import fileinput

	ProbA = [1.0]
	ProbB = [0.0]

	def solve(n,m,N):
		for i in range(N):
			a2b = ProbA[i] * 1/n		# Aから取り出されてBに移る確率
			b2b = ProbB[i] * (1-1/m)	# Bから取り出されず残る確率
			b2a = ProbB[i] * 1/m		# Bから取り出されてAに移る確率
			a2a = ProbA[i] * (1-1/n)	# Aから取り出されず残る確率

			ProbA.append(b2a + a2a)
			ProbB.append(a2b + b2b)

	# main
	if __name__ == "__main__":
		for line in fileinput.input():
			if not line.strip():
				continue

		v = line.split()

		n = float(v[0])
		m = float(v[1])
		N = int(v[2])

		solve(n,m,N)
		print str(ProbA[N])

解説

あまり難しくはありません。高校の確率統計の試験に出そうな問題です。

考え方

2回目以降A、Bに赤玉が入っている確率がどうなるかということが解れば問題は解けます。
常に1個ずつ玉を交換するだけなのでn、mの数は変化しません。
まず、A、Bいずれかの箱に赤玉が入っているとした場合、Aから赤玉が取り出される確率は1/n、Bから取り出される確率は1/mです。そして赤玉が同じ箱に残る場合(つまり、取り出されなかった場合)も考える必要があります。Aに赤玉が残る確率は1-1/n、Bに残る確率は1-1/mです。
なのでi回目の試行で赤玉の位置が確定しているならばi+1回目にAに赤玉がある確率は1/m+(1-1/n)、Bにある確率は1/n+(1-1/m)です。
1回目の試行では0回目(初期状態)はAに赤玉があることが確定しているので前述の通りの確率になります。2回目以降は前回A、Bぞれぞれの箱に赤玉があった確率に移動する場合、移動しない場合の確率を掛ければ良いことになります。

solve()

考え方の通りに実装します。
ProbA、ProbBはi回目にA、Bいずれかの箱に赤玉がある確率を保持します。初期値はProbAが[1]、ProbBが[0]なのは明らかです。

i+1回目の試行時にAに赤玉がある確率は、
 B→Aに移動する場合 :ProbB[i] * 1/m ・・・(a1)
 Aから移動しない場合:ProbA[i] * (1-1/n) ・・・(a2)
i+1回目の試行時にBに赤玉がある確率は、
 A→Bに異動する場合 :ProbA[i] * 1/n ・・・(b1)
 Bから移動しない場合:ProbB[i] * (1-1/m) ・・・(b2)
なので、(a1)+(a2)をProbA、(b1)+(b2)を試行回数までProbBに追加して行けばN回目のA、Bに赤玉がある確率がわかります。

最終的にProbAの最後の要素を印字すれば結果になります。

雑感

問題が公開されたのでやってしまいましたが、この解説を書こうとして問題をみたら「CodeIQ数学ゼミ」に参加希望しない人は遠慮してくださいとのことでした。大変失礼しました。
久しぶりのPythonでしたが結構忘れていました。リストへの要素追加がpush()じゃなかったっけ? とか。