「 今週のお題:均等に分配されるカード」問題解答 ( CodeIQ )

はじめに

CodeIQという所で週次で出題される「今週のお題」シリーズに解答しましたので、ネタバレです。

codeiq.jp

問題

今回は次のような問題でした。

  • 入力m,nに対し、1~mのカードを、数字の合計が等しくなるようにn人に分ける方法が何通りあるか出力する ( ただし人は区別しない )

例えばm,n=8,3 であれば、1,5,6-2,3,7-4,8、1,3,8-2,4,6-5,7、1,2,3,6-4,8-5,7 の3通りなので3を出力する、といった具合です。

解説

次のRubyのコードを提出しました。

f=->m,n{
  sum=m*(m+1)/2
  ave=sum/n
  g,h=
    ->rp,mask{
      rp==1 ? 1 :
              ( x=(1..m).max_by{|i|mask[i]>0?0:i}; h[rp,ave-x,x-1,mask|1<<x] )
    },
    ->rp,rs,u,mask{
      rs==0 ? g[rp-1,mask] :
              (-[u,rs].min..-1).reduce(0){|s,i|mask[-i]>0?s:s+h[rp,rs+i,-i-1,mask|1<<-i]}
    }
  sum%n>0||ave<n ? 0 : g[n,0]
}
p f[*eval(?[+gets+?])]

2種類の関数 g,h の相互再帰によるDPで場合の数を数えるよう実装しました。

  • g: 残りカードからrpで指定された人数で分ける時の場合の数を数える関数
  • h: 残りrp人で、うち1人が ( 既に何枚かカードを割り当てられていて ) 合計残り rs であるときの 場合の数を数える関数

いずれも、パラメータ mask をビットマスクとして使うことにより、選択済みのカードを管理します。

重複カウントを防ぐため、ある人に最初に割り当てられるカードは必ず残りの中で最大の数字のものとし、また、1枚ずつカードを割り当てる時は、必ず前のカードより小さいカードを選んでいくものとします。
この上限は h のパラメータ u によって管理します。

なお、カード全ての合計がそもそも人数で割り切れない場合、あるいは、1人あたりのカードの合計が最大のカード n を下回ってしまう場合は解がないので、事前にチェックして撥ねておきます ( 返り値 0 )

終わりに

相互再帰、なんか妙にカッコイイので個人的には好きです。 なお、「大きい方からケースを潰していこう」というノリで、関数 h で-[u,rs].min..-1というレンジを使った ( 符号を反転して使う ) のですが、あんまり意味なかったような気もします。