Problem 221

221 Alexandrian Integers
Problem 221 - Project Euler
似た問題が前にもあった．それは，n^2の約数が鍵になっていたが…
今回はn^2+1の約数が鍵になる．約数を求めるのは結構大変．
とりあえず，素因数分解を利用してみた．

{-# OPTIONS_GHC -fbang-patterns #-}
import Control.Monad (when, guard)
import Control.Monad.ST (ST, runST)
import Data.Array.ST (STUArray, newArray, readArray, writeArray)
import Data.List (sort, group)

alexandrian :: Integer -> [Integer]
alexandrian n = do
  let u = floor $ fromInteger (div n 4) ** (1/3)
      primes = primesUpTo u
  p <- [1..u]
  let v = min p.floor.sqrt.fromInteger.div n $ 4 * p
  d <- takeWhile (<= v).divisors primes $ p^2 + 1
  let a = p * (p+d) * (p + div (p^2+1) d)
  guard $ a <= n
  return a

main = print.(!!149999).sort.alexandrian $ 2 * 10^15

divisors :: [Integer] -> Integer -> [Integer]
divisors ps = sort.map product.sequence.map (scanl (*) 1).group.factors ps

factors :: [Integer] -> Integer -> [Integer]
factors ps n = f n ps
    where f _ [] = []
          f m (p:ps) | p*p > m = [m]
                     | m `mod` p == 0 = p:f (div m p) (p:ps)
                     | otherwise = f m ps

primesUpTo :: Integer -> [Integer]
primesUpTo n =
    runST $ do
      p <- newArray (1,div n 2) True :: ST s (STUArray s Integer Bool)
      let u = floor.sqrt.fromIntegral $ div n 2
          loop !i = when (i <= u) $ sieve i (2*i*(i+1)) >> loop (i+1)
          sieve !i !j = when (j <= div n 2) $ writeArray p j False >> sieve i (j+2*i+1)
          primes !i !ps | i > div n 2 = return.reverse $ ps
                        | otherwise   = do q <- readArray p i
                                           if q then primes (i+1) (2*i+1:ps)
                                             else primes (i+1) ps
      loop 1
      primes 1 [2]

引数が 2 * 10^15 なのは度重なる実験の結果．ずるいといえば，ずるい．