crystal の AtCoder における pow の使用（健忘録）

はじめに

@Itoi さんの記事を拝見して、RubyにInteger#powがあることを思い出しました。
しかしながら、crystalにはmodするpowが無い様です。
とりあえず、自作関数を作ります。

自作関数(Ruby)

def modinv(n, mod)
  n.pow(mod - 2, mod)
end

組み込み関数がありますので、シンプルです。

自作関数(crystal)

def pow(v, p, mod)
  ret = 1i64
  while p > 0
    if (p & 1i64) != 0
      ret *= v
      ret %= mod
    end
    v *= v
    v %= mod
    p >>= 1
  end
  ret
end
def modinv(n, mod)
  pow(n, mod - 2, mod)
end

自作と言いましても、こちらから拝借したものです。(MIT licenseなので大丈夫！？)

ちなみに、ret = 1i64とInt64型にしておかないと、途中でOverflowErrorが発生します。

C - 経路

Ruby版

def nPk(n, k, mod)
  r = 1
  while k > 0
    r *= n
    r %= mod
    n -= 1
    k -= 1
  end
  r
end
def modinv(n, mod)
  n.pow(mod - 2, mod)
end
w, h = gets.split.map(&:to_i)
mod = 1_000_000_007
ans = nPk(w + h - 2, w + h - 2, mod)
ans *= modinv(nPk(w - 1, w - 1, mod), mod)
ans *= modinv(nPk(h - 1, h - 1, mod), mod)
puts ans % mod

crystal版

def nPk(n, k, mod)
  r = 1i64
  while k > 0
    r *= n
    r %= mod
    n -= 1
    k -= 1
  end
  r
end
def pow(v, p, mod)
  ret = 1i64
  while p > 0
    if (p & 1i64) != 0
      ret *= v
      ret %= mod
    end
    v *= v
    v %= mod
    p >>= 1
  end
  ret
end
def modinv(n, mod)
  pow(n, mod - 2, mod)
end
w, h = read_line.split.map(&.to_i)
mod = 1_000_000_007
ans = nPk(w + h - 2, w + h - 2, mod)
ans *= modinv(nPk(w - 1, w - 1, mod), mod)
ans %= mod
ans *= modinv(nPk(h - 1, h - 1, mod), mod)
puts ans % mod

また途中で、ans %= modを入れないとOverflowErrorが発生します。
ということは、Ruby版でも整数のサイズが32bitから64bitに自動的に変換されていることが推測されます。

	Ruby	crystal
実行時間(ms)	75	10

D - Knight

Ruby版

def nCk(n, k, mod)
  a = 1
  b = 1
  i = 0
  while i < k
    a *= n - i
    a %= mod
    b *= i + 1
    b %= mod
    i += 1
  end
  a * modinv(b, mod) % mod
end
def modinv(n, mod)
  n.pow(mod - 2, mod)
end
x, y = gets.split.map(&:to_i)
mod = 10**9 + 7
if (x + y) % 3 != 0
  puts 0
  exit
end
n = (2 * y - x) / 3
m = (2 * x - y) / 3
if n < 0 || m < 0
  puts 0
  exit
end
puts nCk(n + m, n, mod)

crystal版

def nCk(n, k, mod)
  a = 1i64
  b = 1i64
  i = 0
  while i < k
    a *= n - i
    a %= mod
    b *= i + 1
    b %= mod
    i += 1
  end
  a * modinv(b, mod) % mod        
end
def pow(v, p, mod)
  ret = 1i64
  while p > 0
    if (p & 1i64) != 0
      ret *= v
      ret %= mod
    end
    v *= v
    v %= mod
    p >>= 1
  end
  ret
end
def modinv(n, mod)
  pow(n, mod - 2, mod)
end
x, y = read_line.split.map(&.to_i)
mod = 10**9 + 7
if (x + y) % 3 != 0
  puts 0
  exit
end
n = (2 * y - x) // 3
m = (2 * x - y) // 3
if n < 0 || m < 0
  puts 0
  exit
end
puts nCk(n + m, n, mod)

	Ruby	crystal
実行時間(ms)	87	12

まとめ

• crystalのpow対策に詳しくなりました
• Ruby側のdef modinvをdef pow; def modinvにトランスパイルする仕組みがあると便利ですよねえ
• crystalにオープンクラスってあるのかな