[Haskell-cafe] Zumkeller numbers

[ajb at spamcop.net]

G'day all.

Quoting Richard O'Keefe <ok at cs.otago.ac.nz>:

> These lines of Haskell code find the Zumkeller numbers up to 5000
> in 5 seconds on a 2.2GHz intel Mac.  The equivalent in SML took
> 1.1 seconds.  Note that this just finds whether a suitable
> partition exists; it does not report the partition.

This takes 0.1 seconds on a 2GHz Dell running Linux:

module Main where

import Control.Monad
import qualified Data.List as L

primesUpTo :: Integer -> [Integer]
primesUpTo n
     | n < 13 = takeWhile (<= n) [2,3,5,11]
     | otherwise = takeWhile (<= n) primes
     where
         primes = 2 : 3 : sieve 0 (tail primes) 5
         sieve k (p:ps) x
                 = let fs = take k (tail primes)
                   in [x | x<-[x,x+2..p*p-2], and [x `rem` p /= 0 | p<-fs] ]
                         ++ sieve (k+1) ps (p*p+2)

pseudoPrimesUpTo :: Integer -> [Integer]
pseudoPrimesUpTo n
     | n <= wheelModulus = takeWhile (<= n) smallPrimes
     | otherwise         = smallPrimes ++ pseudoPrimesUpTo' wheelModulus
     where
         largestSmallPrime = 11
         verySmallPrimes = primesUpTo largestSmallPrime
         wheelModulus = product verySmallPrimes
         smallPrimes = primesUpTo wheelModulus

         wheel = mkWheel 1 [1] verySmallPrimes

         mkWheel _ rs [] = rs
         mkWheel n rs (p:ps) = mkWheel (p*n) (do
                                 k <- [0..p-1]
                                 r <- rs
                                 let r' = n*k + r
                                 guard (r' `mod` p /= 0)
                                 return r') ps

         pseudoPrimesUpTo' base
             | n <= base + wheelModulus
                 = takeWhile (<= n) . map (+base) $ wheel
             | otherwise
                 = map (+base) wheel ++ pseudoPrimesUpTo' (base + wheelModulus)

-- Integer square root.
isqrt :: Integer -> Integer
isqrt n
   | n < 0     = error "isqrt"
   | otherwise = isqrt' ((n+1) `div` 2)
   where
     isqrt' s
         | s*s <= n && n < (s+1)*(s+1) = s
         | otherwise                   = isqrt' ((s + (n `div` s)) `div` 2)

-- Compute the prime factorisation of an integer.
factor :: Integer -> [Integer]
factor n
     | n <= 0           = error "factor"
     | n <= primeCutoff = factor' n (primesUpTo primeCutoff)
     | otherwise        = factor' n (pseudoPrimesUpTo (isqrt n))
     where
         primeCutoff = 1000000

         factor' 1 _ = []
         factor' n [] = [n]
         factor' n (p:ps)
             | n `mod` p == 0 = p : factor' (n `div` p) (p:ps)
             | otherwise      = factor' n ps

-- The only function used from above is "factor".

zumkeller :: Integer -> Bool
zumkeller n
     | isPrime            = False
     | isPrime            = False
     | sigma `mod` 2 == 1 = False
     | sigma < 2*n        = False
     | otherwise          = sigmaTest ((sigma - 2*n) `div` 2) (factors n)
     where
         isPrime = case factorn of
                    [_] -> True
                    _   -> False
         factorn = factor n
         sigma = product . map (sum . scanl (*) 1) . L.group $ factorn
         factors n = reverse [ f | f <- [1..n `div` 2], n `mod` f == 0 ]

         sigmaTest 0 _ = True
         sigmaTest _ [] = False
         sigmaTest n (f:fs)
             | f > n     = sigmaTest n fs
             | otherwise = sigmaTest (n-f) fs || sigmaTest n fs

main = print (filter zumkeller [1..5000])

Yes, I cheated by using a different algorithm.

Cheers,
Andrew Bromage