見出し画像

primeSig を使って

競プロ典型 90 問 85 日目 - Multiplication 085(★4) 

import Control.Monad
import Data.Bool
import Data.List

main = readLn @Int >>= print . sol

sol k = (t-3*(d-s)-s) `div` 6 + d
  where
  (t,d,s) = foldl' go (1,1,1) $ primeSig k
  go (t,d,s) i = (
    t*(i+2)*(i+1) `div` 2,
    d*((i+2) `div` 2),
    s*bool 0 1 (i `mod` 3==0))

primeSig :: Integral a => a -> [a]
primeSig = fmap genericLength . group . primeFactors

primeFactors :: Integral a => a -> [a]
primeFactors 1 = []
primeFactors n = factor n primes

factor :: Integral a => a -> [a] -> [a]
factor n (p:ps) 
  | p*p > n   = [n]
  | r == 0    = p:factor q (p:ps)
  | otherwise = factor n ps 
  where (q,r) = quotRem n p

primes :: Integral a => [a]
primes = sm++lg
  where
  1:p:cs = f $ foldl g (1,[1]) sm
  sm = [2,3,5,7]
  lg = p:filter pr cs
  pr n = all ((>0) . mod n) $ takeWhile (\p -> p*p <= n) lg
  f (n,rs) = [n*k+r | k <- [0..], r <- rs]
  g (n,rs) p = (p*n,[r2 | k <- [0..(p-1)], r <- rs, let r2 = n*k+r, mod r2 p>0])

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

karoyakani
ありがとう