HaskellでModintを実装 - Potato's Diary

最近Haskellの勉強を始めました。最初に作りたかったものがModintだったので、型クラスの定義の仕方を学んで早速作ってみました。

Modintとは

競プロの数え上げなどの問題でよく見るこれ

次の条件を満たす○○が何通りあるか求めてください。

(条件1)

(条件2)

…

ただし、答えは非常に大きくなる場合があるので、998244353で割った余りを出力してください。

で使うアレです。modを意識せずに計算できるので、Modintがあると便利です。

実装

class (Integral a) => Modint a where
    p :: a
    modp :: a -> a
    (+%) :: a -> a -> a
    (-%) :: a -> a -> a
    (*%) :: a -> a -> a
    (^%) :: a -> a -> a
    inv :: a -> a
    (/%) :: a -> a -> a

instance Modint Int where
    p = 998244353
    modp a = mod a p
    a +% b = modp $ a+b
    a -% b = modp $ a-b
    a *% b = modp $ a*b
    a ^% 0 = 1
    a ^% x = let aa = modp $ (a ^% (div x 2)) ^ 2 in aa *% (a ^ (mod x 2))
    inv a = a ^% (p-2)
    a /% b = a *% (inv b)

はい。特に解説はしません。

階乗とcombinationも

Modintと一緒に使うことが多い階乗やcombinationもついでに実装します。

import qualified Data.Vector as V

facts :: V.Vector Int
facts = V.scanl (\acc x -> acc *% x) 1 $ V.generate 200200 (+1)
fact :: Int -> Int
fact = (facts V.!)

comb :: Int -> Int -> Int
comb n r
    | n < 0 || r < 0 || n < r = 0
    | otherwise = (fact n) /% (fact r) /% (fact $ n-r)

上のコードでは $200200!$ まで計算できますが、必要に応じて書き換えます。

使用例

この問題を使ってきちんと実装できているか確認します。
atcoder.jp
解説にある通り、以下を計算するだけです。
$\sum_{k=0}^K M \times _{N-1}C_k \times (M-1)^{N-1-k}$

提出コード:
atcoder.jp

終わりに

提出コードの実行時間 (493ms) を見てわかる通り、高速化はできていません (階乗がボトルネック?)。Haskellの知識がついてきたら少しずつ改良したいです。