「 今週のお題:整数倍の得票数」問題解答 ( CodeIQ )

はじめに

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

codeiq.jp

問題

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

  • 空白区切りの入力 m,n に対して、m人の候補者にn人が投票する選挙を考えるとき、以下の条件を満たす得票数パターンが何通りかを出力する。
    • 白票はない
    • 候補者は少なくとも1票以上獲得する
    • 候補者の獲得票数は、最低獲得者の票数の倍数となっている

選挙シーズンということからの問題でしょうか。なお、あくまで「得票数」のパターンであって、誰が票を獲得したかまでは区別しません。つまり、例えば「候補者A: 3票、候補者B: 2票」と「候補者A: 2票、候補者B: 3票」は同じとみなします。

解説

最低得票数を$b$とすると、まず$n$は$b$で割り切れ、かつ$\frac{n}{b}\ge m$です。

そうすると、最低得票数の1人を除いた$m -1$人には、$b$票の塊が$\frac{n}{b}-1$個割り当てられるため、「$\frac{n}{b}-1$を1以上の数$m -1$個に分割する方法」を数えるのと同値です。 分かり易さのため、これを「$\frac{n}{b}-m$を0以上の数$m -1$個に分割」と置き換え、再帰的に計算します。

これをRubyで実装した次のコードを提出しました。

f=->r,k,u{k<2?1:((r+1)/k..u).reduce(0){|s,e|s+f[r-e,k-1,[r-e,e].min]}}
m,n=gets.split.map(&:to_i)
puts (1..n/m).reduce(0){|s,b|n%b>0?s:s+f[t=n/b-m,m-1,t]}

上記説明の再帰的な計算は関数fが担当します。
関数fでは、最大の数を選んでいき、残りをどう分割するかを再帰します。( パラメタ u がその上限 )

これで一応ACしたのですが…。ただ、ズボラな実装なため実行速度が…ということをtweetしたらツッコミを頂きました。

はい。すいませんでした。遅い原因は明白で、メモ化をしなかったからですね。というわけで、簡易ながらメモ化を行った実装が次の通りです。…Rubyでは配列をそのままキーに使えるので便利ですね。

h={}
f=->r,k,u{r<1?1:k<3?u-(r-1)/k:h[[r,k,u]]||=((r+1)/k..u).reduce(0){|s,e|s+f[r-e,k-1,[r-e,e].min]}}
m,n=gets.split.map(&:to_i)
puts (1..n/m).reduce(0){|s,b|n%b>0?s:s+f[t=n/b-m,m-1,t]}

メモ化に加えて、k=2 の場合に再帰せずに直接値を計算することでの高速化も行っています。手元の環境で速度が7倍くらいにはなったので、まあこんなものでしょう。

終わりに

以前は、Rubyで配列を直接ハッシュのキーに使えることを意識していなかったので、自身でキーの計算式を作ってたりしたのですが、直接はやっぱり楽ですね。…メモ化できる時はさぼらずやっていこうと思います。