CodeIQ:どの会場でも均等にセミナーを受講して!

私自身が表題の問題を解いた時のプログラムについて解説します。
問題の詳細は「どの会場でも均等にセミナーを受講して!」(CodeIQ)を参照してください。

問題の概要

何かの展示会のようなものを考えます。
同時間帯にm個のセッションが行われ、nコマあるとします。
この時、すべての会場に同じ回数行くパターンが何通りあるかを求めることにします。
同じ時間帯に複数のセッションに出ることはできませんし、セッションに出ない時間帯があっても構いません。
また、今回は簡単にするため、CLOSEDのセッションや、連続したセッションは考えないものとします。
(つまり、すべての会場でいずれの時間帯もセッションが開催されており、1コマごとに別会場へ移動できます。)

【問題】
標準入力からセッション数 m とコマ数 n がコンマ区切りで与えられたとき、すべての会場に同じ回数行くパターンが何通りあるかを標準出力に出力してください。
(同時に行われるのは最大6個のセッション、コマ数は最大10コマとします。)

私のプログラム

Pythonで解答しています。

#!/usr/local/bin/python3
import fileinput

'''
	組み合わせ数を計算する
'''
def C(n,r):
	ret = 1
	for i in range(n, n-r, -1):
		ret *= i
	for i in range(r, 1, -1):
		ret //= i
	return ret

'''
	@param n コマ数
	@param m セッション数
	@return 組み合わせ数
'''
def takeSession(n,m):
	# nもしくはmが0の場合、どれも行かないとして1としたが、テストパターンでは0らしい
	if (m==0) or (n==0):
		return 0

	# セッション数がコマ数より多い場合、どれも行かない場合だけしかない
	if(m>n):
		return 1

	t = n // m	# 1会場について*最大*何セッション出席すべきか
	ret = 0		# 答え

	# 各会場を最大t回回ることができる
	for i in range(t, 0, -1):
		r = n		# 残りのコマ数
		c = 1		# 各会場をi回まわる場合の組み合わせ数
		for j in range(m):
			c *= C(r,i)
			r -= i
		ret += c

	# 1つも出席しない場合を加算する
	return ret+1

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

		m,n = map(int, line.strip().split(","))

		ret = takeSession(n,m)
		print(ret)

解説

この問題は珍しくプログラムでシミュレーションしたりするよりも、数学的に式を考えて解く方が簡単(というか、先に数学的に解くことができた)問題です。

考え方

具体例を挙げて説明します。次のような会場・時間割になっているとします。
A会場B会場C会場
A-1B-1C-1
A-2B-2C-2
A-3B-3C-3
A-4B-4C-4
A-5B-5C-5
A-6B-6C-6

各会場を同じ回数回るのでA〜C会場をそれぞれ2回、1回、0回回るパターンがあります。
まず、各会場を2回回ることだけを考えます。この場合、A会場を回るパターンは6マスから2マスを選ぶので6C2です。次にB会場を回るのは4C2パターンです。何故ならA会場に行っている時間は他の会場には行けないのでその分選択肢が減ります(例えば、A-1とA-2を回ってしまったらB-1とB-2は選択でません)。同様にC会場は2C2パターンになります。つまり各会場を2回ずつ回るのは6C2×4C2×2C2になります。
同じように1回ずつ回る場合は6C1×5C1×4C1パターンあります。
どこも回らないのは1パターンあります。
よって、(6C2×4C2×2C2) + (6C1×5C1×4C1) + 1がこの場合の答えです。

takeSession()

前述の考え方を一般化して実装したのがこの関数です。
まず、セッションをしていない場合は0を返します。これは問題からは1にすべきか0にすべきかはっきりしなかったのですが0とみなすようです。また、会場数が時間数よりも多い場合は1つもいかないという選択肢しかないので1を返します。

次に、各会場を最大何回回れるのかを計算します(29行目)。これは時間数÷会場数(小数点以下切り捨て)なのは簡単にわかります。最大回数がわかったら1〜最大回数までそれぞれ何通りあるかを求めて足し合わせればOKです。
33行目のループがこの処理にあたります。39行目の足し算は次に説明する各会場を1回、2回、…、回る場合のパターンを足し合わせる処理をしています。

36行目のループは決まった回数(例えば2回ずつ)各会場を回る組み合わせの数を計算します。会場数をn、会場を訪れる回数をxとすると1つ目の会場を訪れるパターンはnCxで、2つ目の会場を訪れる回数はn-xCx、3つ目の会場ならn-2xCx、…で、これを掛け合わせることになります。
このn、n-x、n-2xを計算しているのが38行目で、組み合わせを計算して掛け合わせているのが37行目です。

各会場を1回以上回るパターンの合計に対し、1つも行かない場合の1を加えたのが答えです。

雑感

この問題は数学の問題として解くことができて非常にすっきりしました。