CodeIQ:均等に分配されるカード

私自身が表題の問題を解いた時のプログラムについて解説します。
問題の詳細は「均等に分配されるカード」(CodeIQ)を参照してください。

問題の概要

問題を引用します。
m 枚のカードがあり、それぞれに 1~m までの数字が1つずつ書かれています。
これらのカードを n 人に配りたいと考えています。
(もちろん n < m です。)
それぞれが持つカードの和が、全員同じになるような分け方が何通りあるかを求めてください。
例えば、m = 3, n = 2のとき、1, 2, 3の3枚のカードを「1, 2」と「3」に分ければ、それぞれのカードの和は「3」で一致します。

標準入力から m と n がコンマ区切りで与えられたとき、分け方のパターンが何通りあるかを標準出力に出力してください。
(m は最大で16とします。)
m = 3, n = 2のときは、上記の1通りしかありませんので、以下のように出力します。

(中略)

同様に、m = 7, n = 2のときは以下の4通りがあります。
「1, 2, 4, 7」と「3, 5, 6」に分ける
「1, 2, 5, 6」と「3, 4, 7」に分ける
「1, 3, 4, 6」と「2, 5, 7」に分ける
「1, 6, 7」と「2, 3, 4, 5」に分ける

私のプログラム

Rubyで解答しています。

#!/usr/bin/ruby

$BIT_PATTERN = []
$RESULT = []

# arrに含まれる数値-1のビットをフラグとした整数値を返す
def toBitPattern(arr)
	ret = 0
	for i in arr
		ret |= 1 << (i-1)
	end
	return ret
end

# DEBUG
def printBitPattern()
	arr = $BIT_PATTERN.map{|i| i.to_s(2)}
	p arr
end

# numsを組み合わせてaveに等しくなる値のリストを作成する
# ただし、値は相当するビット位置-1をフラグとした数値
# (例) [1,7]なら0b0000000001000001
# @param nums 値の作成に使える数値のリスト
# @param ave 作成する値の合計値
# @param lst 作成途中の候補値
def makePattern(nums, ave, lst=[])
	for i in nums
		rest = nums.dup	# 使っていない値をコピー
		now = lst.dup		# 現在のリストをコピー

		now << rest.delete(i)

		sum = now.reduce(:+)
		if sum == ave then
			$BIT_PATTERN << toBitPattern(now)
		elsif sum < ave
			rest.delete_if{|j| j <= i}
			makePattern(rest, ave, now)
		end
	end
end

# patternからselと同じ値を含む項目を除外した配列を返す
# @param sel すでに結果候補として選択された値おリスト
# @param pattern 結果構築に使える値の候補
# @return selに含まれる要素と値の重複のない値のリスト
def getRemovedPattern(sel, pattern)
	ret = []
	for r in pattern
		if (sel & r) != 0 then next
		else ret << r
		end
	end
	return ret
end

# 結果を構築する
# $BIT_PATTERNの要素を組み合わせて結果になる値をリストアップする
# @param n 入力値n
# @param fixed 結果作成途中の確定項
# @param pattern 結果を構築するために使える残りのパターン
def makeResult(n, fixed, pattern)
	if fixed.size == n then
		$RESULT << fixed
		return
	end

	if pattern.empty? then return end

	last = pattern.pop
	nf = fixed.dup
	nf << last
	cands = getRemovedPattern(last, pattern)
	makeResult(n, nf, cands)

	makeResult(n, fixed, pattern)
end

# DEBUG
def printResult()
	for arr in $RESULT
		r = arr.map{|i| i.to_s(2)}
		p r
	end
end

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

	m,n = line.split(",").map{|i| i.to_i}

	nums = (1..m).to_a
	total = nums.reduce(:+)

	ave = 0
	if total % n != 0 then
		p 0
	else
		ave = total / n
		makePattern(nums, ave)
#		printBitPattern()
		makeResult(n, [], $BIT_PATTERN)
#		printResult()
		p $RESULT.size
	end

end

解説

★★の問題ですが(私には)かなり難しかったです。
同じ出題者の★★★★の問題よりも難しく感じました。

考え方

プログラムは2段階で処理しています。

その前に、値をいくつにすれば良いのかというのは簡単にわかります。1〜mまでをnで割った値なのは明らかです。これが整数にならない場合、そのような分け方が存在しないことは明らかです。
この値を以降「平均値」と呼びます。

処理の1段階目ではカードが「平均値」になる組み合わせを列挙します。
2段階目で列挙した組み合わせをn組組み合わせて同じカードの重複がないパターンを数えます。このパターンの合計が答えになります。2段階目の処理はRubyの場合combination()を使えば簡単にできますが、時間制限を超えてしまうため他の方法にしなければなりません。

makePattern()

1段階目の処理を行う関数です。
処理としては使用可能な数字のリストから1,2,3,……と取り出して候補リストに追加し、候補リストの合計がaveに等しくなったらパターンとして記録し、合計がave未満なら使っていない値を同じように取り出して候補リストに追加して合計……、という処理を繰り返すことでパターンを列挙するという処理です。
再帰なので見通しが悪いため、具体的に(1..3)の場合で処理を追いかけてみます。

まず、最初引数nums=[1..3]、ave=3が渡ってきます。
32行目の処理でrest=[1..3](numsに等しい)から1を取り出してnowに追加します。1回目はnow=[1]、rest=[2,3]になります。
34行目で合計しますがnow=[1]の合計はave=3未満なので37行目のelsif条件になります。38行目の処理は重複して列挙することを防ぐための処理です。now=[1]の時点では関係ありませんがそれ以外の場合に必要になります(例えば2を取り出した時にrest=[1,3]になっているので1を消す)。
39行目でもう一度自分自身を呼び出しますが、今度はnums=[2,3]、ave=3、lst=[1]で呼び出されます。なので28行目では2、3の順で値が取り出され、now=[1,2]になった時点で35行目の条件に合致してパターンが記録されます。
nums=[1,3]の時は合計がave=3より大きいので無視されます。

パターンを記録する際、toBitPattern()を通してnowの値をビット列に直します。例えば[1,2]は0b0000000000000011になります(1ビット右にずれた状態になる)。
これは組み合わせを行う時に処理を簡単にし、高速化するためです。
パターンのリストは$BIT_PATTERNに記録します。

makeResult()

うまく数値の重複なく組み合わせを列挙し、条件に合うものをカウントします。
手順は次のようになります。
未使用のパターン(最初は$BIT_PATTERNに等しい)から任意のパターン1つを取り出します。その値と重複する数字を持たないパターンのみを抽出し、同じ処理を適用して組み合わせのリストを追加て行きます。
これを組み合わせに含まれるパターンがが入力値nに等しくなるまで繰り返します。
やはり、再帰で見通しが悪いので具体的な値で説明します。

今度はm=7の場合を使用します(m=3だとすべての処理を通らないため)。 m=7,n=2を満たすパターン($BIT_PATTERNに含まれる数の組み合わせ)は次の通りです。(カッコ内の値は実際のビット表現)
[1, 2, 4, 7](0b01001011)
[1, 2, 5, 6](0b00110011)
[1, 3, 4, 6](0b00101101)
[1, 6, 7](0b01100001)
[2, 3, 4, 5](0b00011110)
[2, 5, 7](0b01010010)
[3, 4, 7](0b01001100)
[3, 5, 6](0b00110100)

まず、最初はn=2、fixed=[](空の配列)、pattern=[0b01001011, 0b00110011, 0b00101101, 0b01100001, 0b00011110, 0b01010010, 0b01001100, 0b00110100]で呼び出されます。 71行目のpop()で[3, 5, 6](0b00110100)がpatternから呼び出されます。
72〜73行目でnf=[0b00110100]になり、getRemovedPattern()で残りのパターンからlastに選んだ値と重複する値を持たない値を1つ選びます(後述)。今回の場合、[1, 2, 4, 7](0b01001011)だけが選択されます。cards=[0b01001011]になります。
そして再度makeResult()を呼び出しますが、今度はn=2、fixed=[0b00110100], pattern=[0b01001011]になっているので、同じ処理を適用するとnf=[0b00110100, 0b01001011]、cards=[]となります。
もう一度makeResult()を呼び出し、n=2、fixed=[0b00110100, 0b01001011], pattern=[]なので、64行目の条件に合致し結果を$Resultに記録して終わります。

77行目の呼び出しは選択したものを含まないpatternで呼び出すことで、今回選択したもの以外のもの以外でチェックします。m=7, n=2の場合は同じパターンが別の組み合わせで疲れれることはありませんがm=8, n=3の場合、組み合わせは、
[1, 2, 3, 6], [4, 8], [5, 7]
[1, 3, 8], [2, 4, 6], [5, 7]
[1, 5, 6], [2, 3, 7], [4, 8]
となり、[4,8]と[5,7]は2回使われます。
77行目の呼び出しはこのような場合に対応するために必要な処理になります。

getRemovedPattern()

前述しましたがpatternから引数selに選ばれた値と重複する番号を持つパターンを削除するための処理です。これによって繰り返しごとのパターン数が大幅に減少してゆくので計算量を大きく削減できます。
パターンごとに含まれる番号はビット表現になっているので、重複がなければ論理積を取ると0になります。なので、論理積が0になった値だけを選択して返します。
これをビット演算にせず配列の差集合を求めるなどすると計算量がかなり増えてしまいます。

雑感

先に書いた通りかなり苦労しました。
半分(パターンを列挙する部分)はすぐにできたのですがそこからが遠かった。組み合わせを列挙するうまい方法が考えつかず、別のロジック(パターンを列挙してから組み合わせを作るのではなく、いきなり組み合わせを求めてしまう)なども考えたのですが速度性能を満たせないで苦労しました。
多分、もっと効率の良い(一つの関数で解決できるような)ロジックがあると思います。