見出し画像

infix 関数適用演算子が早い - 素数で割った余りの出力

競プロ典型 90 問 15 日目 - Don't be too close(★6)

{-# LANGUAGE BangPatterns #-}

import Control.Monad
import Data.List
import Data.Vector.Unboxed ((!))
import qualified Data.Vector.Unboxed as U

main = readLn >>= mapM_ print . U.toList . sol

sol n = U.map (f n) $ U.generate n succ
  where
  f n k = foldl' (.+.) 0 $ map (\i -> comb (n-(k-1)*(i-1)) i) [1..(n+k-1) `div` k]
  comb n m = fac!n .*. inv!m .*. inv!(n-m)
  fac = U.constructN (n+1) g
  inv = U.constructN (n+1) h

g u
  | U.null u  = 1
  | otherwise = U.last u .*. U.length u

h u
  | U.null u      = 1
  | U.length u==1 = 1
  | otherwise     = U.last u .*. pwr (U.length u) (p-2)

p = 10^9+7 :: Int

(.*.) = ((`mod` p) .) . (*)
infixl 7 .*.

(.+.) = ((`mod` p) .) . (+)
infixl 6 .+.

pwr !a !n
  | n == 0    = 1
  | odd n     = a .*. r
  | otherwise = r
  where
  q = pwr a (n `div` 2)
  r = q .*. q

いいなと思ったら応援しよう!

karoyakani
ありがとう