HaskellでModIntを実装してみた - ゆきくらげのブログ

HaskellでModIntを実装してみた(簡単な解説付き)

注意

・あまり詳細な解説はついていません．
・わからないところがあって，わからないと書いてあります．今後理解したら追記するかもしれません．
・ぬくぬく理解でやっているので間違いが含まれているかもしれません．指摘してもらえるとありがたいです．

動機

この問題を見てください．この後普通にネタバレします．
atcoder.jp
この問題の答えは

$(P-1)(P-2)^{N-1}\ mod\ 1000000007$
です．ここは解答を見てもらえばわかります．
ここで制約を見てみましょう．
$N<10^9$
とあります．したがって素直に書くと(p-2を10^9回かけることになるので)TLEしますし，オーバーフローもします．したがって……

ModIntを実装して，繰り返し2乗法をできるようにしよう！

が今回の動機です．

実装

先に実装を書いておきます．

{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Main where

import qualified Data.Proxy   as Proxy
import qualified GHC.TypeNats as TypeNats
import           Prelude      hiding (negate, recip, (*), (+), (-), (/), (^),
                               (^^))
import qualified Prelude

main :: IO ()
main = return ()

------------
-- ModInt --
------------

newtype Mod a (p :: TypeNats.Nat) = ModInt a deriving (Eq, Show)

instance (TypeNats.KnownNat p, Integral a) => Num (Mod a p) where
  (ModInt x) + (ModInt y) = ModInt $ (x + y) `mod` p where
    p = fromIntegral $ TypeNats.natVal (Proxy.Proxy :: Proxy.Proxy p)
  (ModInt x) * (ModInt y) = ModInt $ (x * y) `mod` p where
    p = fromIntegral $ TypeNats.natVal (Proxy.Proxy :: Proxy.Proxy p)
  negate (ModInt x) = ModInt $ - x `mod` p where
    p = fromIntegral $ TypeNats.natVal (Proxy.Proxy :: Proxy.Proxy p)
  abs = id
  signum _ = 1
  fromInteger n = ModInt $ fromInteger n `mod` p where
    p = fromIntegral $ TypeNats.natVal (Proxy.Proxy :: Proxy.Proxy p)

instance (TypeNats.KnownNat p, Integral a) => Fractional (Mod a p) where
  recip (ModInt n)
    | gcd n p /= 1 =
      error "recip :: Mod a p -> Mod a p : The inverse element does not exist."
    | otherwise = ModInt . fst $ extendedEuc n (-p)
    where
    p = fromIntegral $
      TypeNats.natVal (Proxy.Proxy :: Proxy.Proxy p)
  fromRational r = ModInt n / ModInt d where
    n = fromInteger $ Ratio.numerator r
    d = fromInteger $ Ratio.denominator r

解説

*** 環(Ring)と体(Field)
　まず環と体を実装しました．こちらの説明が分かりやすいです．ありがとうございます.
zellij.hatenablog.com
　環には加法と乗法の演算(+と*)，加法の単位元(zero)，加法の逆元を求める演算(negate)を実装し，体はそれに加えて乗法の単位元(one)と逆元を求める演算(recip)を実装しました．HaskellのNumクラスは四則演算に加えて，絶対値(abs)の定義が必要であるので，ModIntでは扱いづらく，環と体から作りました．モジュラ計算は法が素数の時に体となるので，そのように実装しました．使うときは素数を入れるようにしてください
　また，通常の四則演算はNumの下で定義されていますが，先述の通りNumの使い勝手が悪いのでデフォルトの四則演算を隠して再実装しています。

2021/09/27Numで実装し直しました(競プロでの使いやすさを重視)

型レベル自然数

　1がInt型，"Hello"がString型，またInt型の変数を受け取ってInt型の変数を返す関数はInt->Int型であるのと同じように，「型の型」のような存在，カインドがあります．IntやStringは具体型と呼ばれカインドでは" * "と表されます．さらに型レベル関数というものもあり，例えばMaybeは具体型aを受け取って具体型Maybe aを返す型レベル関数として解釈できます．よってMaybeのカインドは*->*となります．
　ここで，最終的な目標はModInt n (nは自然数)を具体型にしたいことです．つまり，ModIntは型レベル関数であって，自然数をとって具体型を返すということです．型レベル関数には変数を突っ込むことはできないので，必然的に型レベル自然数が必要となります．しかし，型レベル自然数そのものが具体型にはなって欲しくない(任意の自然数nに対して具体型nができると，意味が分からない)ので，型レベル自然数のカインドは"Nat"とされています．したがって，ModIntのカインドはNat -> * です．
　型レベル自然数を変数に持ち上げるにはnatValを使います．実装を見れば大体わかると思います．所謂「コンパイラマジック」を使っているのでこれで出来るんだね～という理解です．今のところ．
　この際，拡張 DataKinds(デフォルトで"1"とか"2"とかは変数レベルのものを指すのでそれを型レベルにするために必要) と KindSignatures(クラス定義の時カインド注釈をつけれる, ::Natの部分)，ScopedTypeVariables(関数の型注釈で使った型変数をスコープ内に持ち込める．natVal使ってる部分)が必要です．

繰り返し二乗法

　これは簡単なアルゴリズムで
$a^1=a$
$a^2=a*a$
$a^{2n}=(a^n)^2$
$a^{2n+1}=a*(a^n)^2$
を実装すれば大丈夫です。計算量はO(logN)になります

追記

2021-04-23 繰り返し二乗法の説明を追加
2021-04-24 少し修正
2021-05-05 かなり修正
2021-09-27 かなり修正