【Ruby】二項係数(nCr)の素数modを求めるコード

結論

Ruby

# r<=0の挙動には注意
def nCrMod(n, r, mod)
  numerator = ((n - r + 1)..n).inject { |result, item| result * item % mod } # 分子
  denominator = (1..r).inject { |result, item| result * item % mod } # 分母
  puts numerator * (denominator.pow((mod - 2), mod))
end

急ぐ方が多いと思うので始めに持ってきました。
このコードに至る経緯や解説は下に書いていきます。

はじめに

はじめまして、糸井と申します。
2022年4月に友人と競プロを始めてハマっています！
言語はRubyです。競プロで高みを目指すということも目標ですが、Rubyの経験を積むためにC++ではなくRubyを使っています。（競プロで高みを目指すだけなら、間違いなくC++を覚えた方がいいですね）

今回はAtCoderでたまに見るnCrの求め方と、それの素数のmodを求めたので記録しておこうと思います。

実際の問題

Knight [AtCoder Beginner Contest 145 D]
https://atcoder.jp/contests/abc145/tasks/abc145_d

問題文
二次元グリッドの原点 (0,0) にチェスのナイトの駒があります。
ナイトの駒はマス (i,j) にあるとき (i+1,j+2) か (i+2,j+1) の
どちらかのマスにのみ動かすことができます。
ナイトの駒をマス (X,Y) まで移動させる方法は何通りありますか？
10^9+7 で割った余りを求めてください。


制約
1≤X≤10^6
1≤Y≤10^6
入力中のすべての値は整数である。

ざっくりと解釈すると、"上に進んだ回数"と"右に進んだ回数"の二項係数問題になります。

10^6クラスの二項係数の求め方

壁となるのは計算時間です。愚直に計算するとTLEになりました。
解答はC++で書かれており、QiitaなどでRubyで書かれているものもコードを解読し切れませんでした。
その中で、なんとか真似できる解答を高級言語で書かれている記事を見つけました。

Pythonで逆元を使ってnCr mod 1000000007を計算　（コッコさん）
https://cocoinit23.com/ncr-mod-1000000007/

正直に言うと、逆元の定義や求め方、使い方に関しては完璧には理解できていません。かろうじて「かけ算の途中でmodを積極的に求めて良い」ということは理解しました。

コッコさんのPythonのコードからRubyに書き換えてみました。

Ruby

def nCrMod(n, r, mod)
  numerator = ((n - r + 1)..n).inject { |result, item| result * item % mod } # 分子
  denominator = (1..r).inject { |result, item| result * item % mod } # 分母
  puts numerator * (denominator.pow((mod - 2), mod))
end

これをコピペしてn, r, modを渡せばそのまま使えます。

「injectとpowって何だっけ？」という方はメモをどうぞ。

injectについて

Ruby 3.1 リファレンスマニュアル

https://docs.ruby-lang.org/ja/latest/method/Enumerable/i/inject.html
list.injectはlistのたたみこみ演算をしてくれます。

p [2, 3, 4, 5].inject {|result, item| result + item }        #=> 14

リファレンスの例では2+3+4+5をしていますね。
　

powについて

Ruby 3.1 リファレンスマニュアル
https://docs.ruby-lang.org/ja/latest/method/Integer/i/=2a=2a.html
powはint.pow(x, y)でintをｘ乗し、計算途中でmod yを求めてくれます。

解答

上記を使って解いたのがこちらです。

Ruby

def knights()
  arr = gets.split(" ")
  x = arr[0].to_i
  y = arr[1].to_i

  if (x + y) % 3 != 0
    puts 0
  else
    # nCrMod
    n = (x + y) / 3
    r = x - n
    if r < 0 || y - n < 0　# rが負のとき(1..r)で実行時エラー。実際の解は0
      puts 0
    elsif r == 0 # r = 0 のとき同様にエラー。実際の解は1(nC0)
      puts 1
    else
      mod = 10 ** 9 + 7
      numerator = ((n - r + 1)..n).inject { |result, item| result * item % mod }
      denominator = (1..r).inject { |result, item| result * item % mod }
      puts numerator * (denominator.pow((mod - 2), mod)) % mod
    end
  end
end

knights()

nCr mod問題で詰まっているRuby使いさんに届けば幸いです。