CodeIQ:幹事が楽な歓送迎会

私自身が表題の問題を解いた時のプログラムについて解説します。
問題の詳細は「幹事が楽な歓送迎会」(CodeIQ)を参照してください。

問題の概要

問題を引用します。
人事異動が多く発生する時期になりました。
歓送迎会などの飲み会などを開催したとき、幹事さんの悩みの種となるのがお釣りの準備です。
お釣りが出ないように参加者全員がちょうどの金額を用意してくれると助かるのですが、なかなかうまくいかないものです。

そこで、お釣りが不足しないような順番で会費を集めようと思っています。
例えば、会費が4000円のとき、千円札を出す人が1人いれば、五千円札を出す人が4人いても、先に千円札の1人から受け付けると、残りの4人のためのお釣りを準備する必要がありません。

ここでは、会費が3000円のとき、千円札を出す人が m 人、五千円札を出す人が n 人いるものとします。
このとき、事前にお釣りを準備しなくても受け付けられる順番が何通りあるかを求めます。
ただし、出すお札の種類での順番についてのみ考え、どの人が出したかは問わないものとします。

例えば、m = 3, n = 2 のとき、以下の5通りがあります。
(1) 千円札 → 千円札 → 千円札 → 五千円札 → 五千円札
(2) 千円札 → 千円札 → 五千円札 → 千円札 → 五千円札
(3) 千円札 → 千円札 → 五千円札 → 五千円札 → 千円札
(4) 千円札 → 五千円札 → 千円札 → 千円札 → 五千円札
(5) 千円札 → 五千円札 → 千円札 → 五千円札 → 千円札

標準入力から参加者数が与えられるとき、m, nのすべての組み合わせについて、事前にお釣りを準備しなくても受け付けられる順番が何通りあるかを求め、その和を標準出力に出力してください。
なお、参加者は「千円札を出す人」と「五千円札を出す人」のいずれかであるものとし、最大で32人までとします。

例えば、参加者数が5人のとき、m, n の組み合わせと、受け付けられる順番のパターン数は以下の表のようになるため、12を出力します。
mnパターン数
050通り
140通り
232通り
325通り
414通り
501通り
合計12通り

【入出力サンプル】
標準入力
5

標準出力
12

私のプログラム

Rubyで解答しています。

#!/usr/bin/ruby

$Memo = {}

def take(m,n,r)
	key = m.to_s + "," + n.to_s + "," + r.to_s
	if $Memo[key] != nil then return $Memo[key] end

	if r < 0 then return 0 end
	if m < 0 then return 0 end
	if n < 0 then return 0 end
	if (m == 0) && (n == 0) then return 1 end

	ret = 0
	ret += take(m-1, n, r+3)
	ret += take(m, n-1, r-2)

	$Memo[key] = ret

	return ret
end

def solve(x)
	ret = 0

	for n in 0..x
		m = x-n

		if n == 0 then ret += 1
		elsif 3*m < 2*n then ret += 0
		else ret += take(m,n,0)
		end
	end

	return ret
end

# main
while line = gets
	line.strip!
	if line.empty? then next end

	p solve(line.to_i)
end

解説

問題をざっと読んでメモ化再帰かな、と思う問題です。

考え方

メモ化のことは置いておいて再帰でパターンを割り出せることを説明します。
必要な情報は会費未徴収の千円札の人の人数、五千円札の人の人数、千円札の枚数です。
関数が呼ばれるごとに千円札の人数か五千円札の人数を1減らします。
千円札の人数を減らした場合は千円の枚数を+3し、五千円札の人数を減らした場合は千円札の枚数を-2します。
その結果、千円札の枚数がマイナスになった(つまりお釣りが足りなくなった)らNGです。千円札の数が0以上の間に人数が0になったらOKです。後、千円か五千円の人数がマイナスになった場合も(そんなことはできないので)NGです。
これが基本的なロジックです。

で、これをメモ化するのですが千円の人数、五千円の人数、千円の枚数が同じ状態が現れたらその後のパターン数は2回計算する必要はないので、千円の人数、五千円の人数、千円の枚数をキー、パターン数を値としてメモを作成すればOKです。

main

入力値を数値にしてsolve()に渡し、結果を印字します。

solve(x)

千円札の人数と五千円札の人数の全パターンをループでtake()に渡し、その返却値を加算します。これが答えになります。take()は千円札と五千円札の人数が決まっていたらその場合のパターン数を求める関数です。初期値は千円札の人数m、五千円札の人数n、千円札の枚数0枚になります。

私の実装では計算しなくて良い場合を条件文で弾いています。
まず、千円札の人数の3倍が五千円札の人数の2倍未満の場合、必ず千円札が不足するのでパターン数は0です。
全員が千円札の場合は1通りしかないのでこれも定数1を結果に加えます。

take(m,n,r)

引数mは千円の人数、nは五千円の人数、rは千円札の枚数です。

最初にメモに計算結果があるかをチェックします。メモに結果があればそれを返却して終わります。多少早くなるかと思ってメモのキーを文字列にしていますが、配列でも問題ないと思います。

次に終了条件で、次の通りです。
r<0(千円札が足りなくなった)の場合はダメなので0を返す。
m<0(千円札を出す人が足りなくなった)場合もダメなので0を返す。
n<0(五千円札を出す人が足りなくなった)場合もダメなので0を返す。
m==0かつn==0(ちょうど全員から徴収した)場合だけがOKなので1を返す。

それ以外の場合は未徴収の人がいるので徴収するためtake()を再帰呼び出しします。
take(m-1, n, r+3)は千円札の人から徴収した場合なので、mを1減らしてrを+3して再帰呼び出しします。
take(m, n-1, r-2)は五千円札の人から徴収した場合なので、nを1減らしてrを-2して再帰呼び出しします。

関数が戻る前にメモにその[m,n,r]の場合のパターン数を記録します。
結果を返却して終わりです。

雑感

最初に書いた通り、メモ化再帰だなと思ったのですが、やっているときに[m,n,r]の組みでメモに十分ヒットするのか(性能改善に繋がるほど同じパターンが現れるのか)? と一瞬不安になりましたが問題ありませんでした。