[Haskell-cafe] Zumkeller numbers

[Daniel Fischer]

Am Dienstag 08 Dezember 2009 01:54:12 schrieb 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:
<snip>
>
> Yes, I cheated by using a different algorithm.
>
> Cheers,
> Andrew Bromage

Yes, what I do (after fixing an embarrassing typo from the first post). Using Int, it's a 
little faster, with Integer the same time. I've now added the printout of one partition.

module Main (main) where

import Data.List (sort)
import Control.Monad (liftM2)
import System.Environment (getArgs)

main = do
    args <- getArgs
    let bd = case args of
                (a:_) -> read a
                _ -> 4000
    mapM_ (\n -> case zumParts n of
                   [] -> return ()
                   ((l,r):_) -> putStrLn (show n ++ "   " ++ show l ++ "   " ++ show r)) 
[1 .. bd]

trialDivPrime :: [Int] -> Int -> Bool
trialDivPrime (p:ps) n = (p*p > n) || (n `mod` p /= 0) && trialDivPrime ps n

oprs = 5:7:filter (trialDivPrime oprs) (scanl (+) 11 $ cycle [2,4])

primes :: [Int]
primes = 2:3:oprs

factors :: Int -> [(Int,Int)]
factors n
    | n < 0     = factors (-n)
    | n == 0    = [(0,1)]
    | otherwise = go n primes
      where
        go 1 _ = []
        go m (p:ps)
            | p*p > m   = [(m,1)]
            | otherwise =
              case m `quotRem` p of
                (q,0) -> case countFactors q p of
                            (m',k) -> (p,k+1):go m' ps
                _ -> go m ps

countFactors :: Int -> Int -> (Int,Int)
countFactors n p = go n 0
      where
        go m k = case m `quotRem` p of
                    (q,0) -> go q (k+1)
                    _ -> (m,k)

divisors :: Int -> [Int]
divisors n = sort $ foldr (liftM2 (*)) [1] ppds
      where
        ppds = map (\(p,k) -> take (k+1) $ iterate (*p) 1) $ factors n

partition :: Int -> [Int] -> [[Int]]
partition 0 _   = [[]]
partition n []  = []
partition n (p:ps)
    | n < p     = partition n ps
    | otherwise = [p:parts | parts <- partition (n-p) ps] ++ partition n ps

zumParts :: Int -> [([Int],[Int])]
zumParts n
    | sigma < 2*n   = []
    | odd sigma     = []
    | otherwise     = [(prt ++ [n],init divs `without` prt) | prt <- prts]
      where
        divs = divisors n
        sigma = sum divs
        target = (sigma-2*n) `quot` 2
        sml = takeWhile (<= target) $ init divs
        prts = map reverse $ partition target (reverse sml)
        xxs@(x:xs) `without` yys@(y:ys)
            | x < y     = x:xs `without` yys
            | x == y    = xs `without` ys
            | otherwise = xxs `without` ys
        xs `without` _  = xs