競プロ学習記録#2_ABC156復習 逆元

かんだむ

はじめに

↓前回の記事↓　　　　　　　　　　　　　　　　　　　　　　　　　　　　競プロ学習記録#1_Rubyのeachが遅いなんて知らなかったよ……

ABC156に出場したら散々たる成績でした……

敗因は曖昧なままにしていた逆元に見事討ち取られた、ということでしっかり復習していきたいと思います。


ABC156-B Digits

無題

NをKで割り続けて回数を数えるだけの問題、なんだけど問題文読み違えてK進数表記に直すもんだと思ったりして時間使いすぎたよね……

n,k = gets.chomp.split.map(&:to_i)
cnt = 0
loop do
 break if k ** cnt > n
 cnt += 1
end
puts cnt

なんとなくwhileよりloopのほうが書きやすくってloop使ってしまう、前回の記事で触れたブログ記事にloopはめっちゃ遅いよって書いてあったけど……

ただこの問題は折角Ruby使っているのでこう書けばよかったなって思った。

n,k = gets.chomp.split.map(&:to_i)
puts n.to_s(k).size

桁数はto_s.sizeでサクッと求まるし、ついでに引数渡すとK進数にしてくれる、Rubyくん最高！

ABC156-D Bouquet

無題

N個のものから1個以上選んで組を作る、ただしa,b個の組は除く、全部で組み合わせは何通りでしょうかっていう問題。

すべての組み合わせは2^N - 1通り、a,b個の組はそれぞれnCa,nCb通りなので答えは(2^N - 1 - nCa - nCb) % 10^9 + 7 通りとなる。

ここまではコンテスト中にも思いついたが、じゃあ具体的にどう計算を実装していくのかがてんでわからず詰みました。

結局やることは全部以下の神記事に書いてある。散々読んでいるはずなのに……

まず全ての組み合わせ2^N-1通りについて、今回の制約はN = 10^9なので愚直に計算するのは不可能ですが、繰り返し2乗法（二分累乗法）を用いるとO(log N)で高速に求められる、これを知らなかった……

こんな感じに書いてみました。

def mod_pow(base,n,mod)
 if n == 1
   res = base
 elsif n % 2 == 0
   res = mod_pow(base,n/2,mod) ** 2 % mod
 else
   res = mod_pow(base,n-1,mod) * base % mod
 end
 res
end

次に問題なのが二項係数について、記事に乗っているライブラリは今回の場合、Nが大きすぎて使用できません。

ちゃんと理解しないでライブラリでどうにかしようとするからこういうことになるのだ。ちゃんとこれ読んで二項係数理解しようね。

もう今回のケースもバッチリ書いてあるとおり！組み合わせを手計算しやすい形に直すとこうなるわけですが、

画像3

分子はnから下っていってn-k+1までのk回の掛け算、分母はkの階乗でk回の掛け算、ということで十分時間内で計算終わる！僕は本番中は分子の計算量をn-k回の掛け算だと勘違いしてて先へ進めませんでした、あのさぁ……

ということでやりたいことはmodをとった分子にmodをとった分母の逆元をかけてmodをとる！以上！

require 'openssl'
def combination_mod(n,k,mod)
 if n < k || (n < 0 || k < 0)
   0
 elsif k == 0 || k == n
   1
 else
   k = n - k if k > n - k
   x = (n-k+1..n).inject{|res,i| res * i % mod}
   y = (1..k).inject{|res,i| res * i % mod}
   inv_y = OpenSSL::BN.new(y).mod_inverse(mod)
   x * inv_y % mod
 end
end

逆元を求めるのはFermatの小定理や拡張Euclidでちゃんと実装したほうがいい気もするが、OpenSSL::BN.new(y).mod_inverse(mod)が使えるということなので使ってしまった。コードはスッキリするが記事を書いててなんだかまずい気もしてきた、後でちゃんと拡張Euclid実装するかんね……！

まとめ

modの逆元と二項係数、繰り返し2乗法、ぜってえ忘れねえかんな……！

この記事が気に入ったらサポートをしてみませんか？