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])
いいなと思ったら応援しよう!
ありがとう