AtCoder ABC Q110の問題DをHaskellで解いてみました

問題

正整数 N,M が与えられます。

a1×a2×…×aN=M となる正整数からなる長さ N の数列 a が何通りあるかを 10^9+7 で割った余りを求めてください。

ただし、数列 a' と a'' が異なるとは、ある i が存在して ai'≠ai'' であることをいいます。

制約
入力はすべて整数である
1≤N≤10^5
1≤M≤10^9

解法

まず入力Mを素因数分解し2^a+3^b+5^c+,,といった形で表せるようにします。それから長さNの数列となるので求めたa,b,c,,のそれぞれをN-1の仕切りで分けておき、その結果を掛け合わせたものの10⁹+7の剰余が答えとなります。

例えばNが3でMが12の場合,M=2²+3¹と表せるので
2,1をそれぞれ2個の仕切りで分けた場合の組み合わせ数を求めます。
2を2個の仕切りで分けた場合は、
4!/(2!×2!) = 6
1を2個の仕切りで分けた場合は、
3!/(2!×1!) = 3 となるので答えは 6 × 3 = 18となります。

標準入力

Haskellで標準入力から2つの数値を受け取る場合、以下のように記述できます。

(n:m:_)<-L.map read . words<$>getLine::IO[Integer]

素数取得

素数の計算について、5以降の素数は6の倍数の前後に現れるという性質を利用して以下のように求めます。

primes :: Integral a => [a]
primes = L.map fromIntegral ([2, 3] ++ primes' :: [Int])
    where
      primes' = 5 : sieve [] primes' 7
      sieve divs (x : xs) n = ps ++ sieve (divs ++ [x]) xs (x * x)
          where
            isPrime m = and [rem m x /= 0 | x <- divs]
            ns = [y + z | y <- [n, n + 6 .. x * x - 2], z <- [0, 4]]
            ps = L.filter isPrime ns

isPrimeの判定についてx <- divsとしてそれまでの素数全てで割れるか見ていますが、判定対象の数値の平方根分まで見れば良いのでx < sqrt mの条件を追加しても良いのかもしれません。

素因数分解

Haskellでの素因数分解の実装は以下のようになります。

fac :: [Integer]->M.IntMap Integer->Integer->M.IntMap Integer
fac _ m 1 = m
fac [] m n = M.insert (fromIntegral n) 1 m
fac pa@(p:pr) m n
    |mod n (fromIntegral p)==0　=　fac pa (M.insertWith(+) (fromIntegral p) 1 m) (div n (fromIntegral p))
    |otherwise =fac pr m n

これを呼び出すときは以下のようにします。

fac primes M.empty m

第一引数に素数を、第二引数にmapをそして第三引数に素因数分解の対象となる数値を渡します。関数内では先頭の素数から順に取り出して対象の数値を割れるだけ割って、その度にmapをカウントアップさせる、これを対象の数値が1になるまで繰り返すことで素因数分解を実装しています。

組み合わせの数

組み合わせ数計算結果の剰余計算を高速化させるにはモジュラ逆数を使った方法があります。

加算、減算、乗算に比べて徐算は計算コストがかかるのですがモジュラ逆数を使うと組み合わせ数計算結果の剰余計算を徐算を使わずに行うことができるので高速化することができます。

60 / 4 の結果の mod 11について通常であれば、60 / 4の結果を求めて15 mod 11を求めて4というようになります。
4 の徐算結果のmod11についてモジュラ逆数は3になるので、3掛けた結果のmod11になります。
そのため60 × 3 mod 11 = 180 mod 11 = 4となり結果は同じになります。

モジュラ逆数は拡張ユークリッド互除法を使うと求めることができます。
ax+by=gcd(a,b)に対してa=11 b=3を代入する
x=1, y=0の11×1 + 3×0 = 11と
x=0, y=1の11×0 + 3×1 = 3を比較しいく

11/3=3よりx=1, y=0の11×1 + 3*0 = 11を以下のように変換する。
x=1-(0×3)=1 y=0-(1×3)=-3 gcd=11-(3×3)=2
これで現在の計算結果は以下のように表せられる
x=1 y=-3 11×1 + 3×(-3) = 2
x=0, y=1の11×0 + 3×1 = 3

3/2=1よりx=0, y=1の11×0 + 3×1 = 3を以下のように変換する。
x=0-1×1=-1 y=1-(-3)=4 gcd=3-2×1=1
これで現在の計算結果は以下のように表せられる
x=1 y=-3 11×1 + 3×-3 = 2
x=-1, y=4の11×(-1) + 3×4 = 1
2 mod 1 = 0となり、ここの3*4の4の部分がモジュラ逆数となる。

モジュラ逆数を求める計算をHaskellで実装すると以下のようになる。

invmod p g = invmod' 1 0 p 0 1 g
    where invmod' px py pgcd x y gcd
            | m==0 = mod (x+g) g
            | True = invmod' x y gcd (px - (d*x)) (py - (d*y)) m
            where
                m = mod pgcd gcd
                d = div pgcd gcd

最終的な実装

これらを合わせて最終的な実装は以下のようになりました。

import Control.Monad
import Data.List as L
import Data.IntMap as M

primes :: Integral a => [a]
primes = L.map fromIntegral ([2, 3] ++ primes' :: [Int])
    where
      primes' = 5 : sieve [] primes' 7
      sieve divs (x : xs) n = ps ++ sieve (divs ++ [x]) xs (x * x)
          where
            isPrime m = and [rem m x /= 0 | x <- divs]
            ns = [y + z | y <- [n, n + 6 .. x * x - 2], z <- [0, 4]]
            ps = L.filter isPrime ns


fac :: [Integer]->M.IntMap Integer->Integer->M.IntMap Integer
fac _ m 1 = m
fac [] m n = M.insert (fromIntegral n) 1 m
fac pa@(p:pr) m n
    |mod n (fromIntegral p)==0 = fac pa (M.insertWith(+) (fromIntegral p) 1 m) (div n (fromIntegral p))
    |otherwise =fac pr m n


invmod p g = invmod' 1 0 p 0 1 g
    where invmod' px py pgcd x y gcd
            | m==0 = mod (x+g) g
            | True = invmod' x y gcd (px - (d*x)) (py - (d*y)) m
            where
                m = mod pgcd gcd
                d = div pgcd gcd

divmod a b p = mod (a * invmod b p) p

cmbmod n k p = L.foldl' (\a i->divmod (mod (a*(n-i+1)) p) i p) 1 [1..k]

sol 1 _ = 1
sol _ 1 = 1
sol n m = L.foldl' (\a b -> (a * cmbmod (b+n-1) b 1000000007) `mod` 1000000007) 1 fm
    where fm = elems $ fac primes M.empty m

main=do
    (n:m:_)<-L.map read . words<$>getLine::IO[Integer]
    print $ sol n m

新しい言語の学習にAtCoderが有効そうだと思って試して見ましたが、モジュラ逆数など数学の部分で躓きました。数学的の知識などが必要になるのでそれに慣れていけば一度得意な言語で解いて、その後別の学習中の言語で解き直すなどして言語の理解に役立ちそうだと思いました。