Project Euler problem 92 (sequences formed by the sum of squares of digits) in Haskell

Question

I made the following program in Haskell to solve Project Euler's problem 92. The program works (i.e. the solution it yields is correct), but it's not exactly the fastest in the world.

The problem is:

A number chain is created by continuously adding the square of the digits in a number to form a new number until it has been seen before.

For example,

• 44 → 32 → 13 → 10 → 1 → 1
• 85 → 89 → 145 → 42 → 20 → 4 → 16 → 37 → 58 → 89

Therefore any chain that arrives at 1 or 89 will become stuck in an endless loop. What is most amazing is that EVERY starting number will eventually arrive at 1 or 89.

How many starting numbers below ten million will arrive at 89?

My code:

digits :: Int -> [Int]
digits 0 = []
digits x = digits (x `div` 10) ++ [x `mod` 10]

squareDigits :: [Int] -> [Int]
squareDigits x =  map (^2) x

sumSquareDigits :: Int -> Int
sumSquareDigits x = sum (squareDigits (digits x))

chainTo89 :: Int -> Bool
chainTo89 89 = True
chainTo89 1 = False
chainTo89 x = chainTo89 (sumSquareDigits x)

main :: IO()
main = print $ length $ filter (chainTo89) $ [1..9999999]

Any feedback on how this can improved, ranging from advice on good practices to very concrete algorithmic improvements, preferably without using too complicated Haskell concepts (I've only been programming in Haskell for three days) is more than welcome. If you are going to propose rather complex improvements (e.g. state monads?), then that's fine, so long as it's understandable to a beginner. :)

Note: I've implemented this in Python too, where I mapped known outcomes to a dictionary, which sped up the process considerably (i.e. the moment the chain function encounters a number already checked, it goes straight to the outcome without chaining further). However, I don't know how to do that in Haskell, if that is good practice in Haskell in the first place. This is exactly one of the reasons I asked this question in the first place; it feels as if it would be good to have some 'global' variable that can be used to map known values to, but that seems to go against Haskell's spirit.

Answer 1

Overall, your solution looks like a reasonable Haskell implementation of the question. Well done. I found some tiny bits to comment about anyway.

Separate algorithm from IO

Part of your algorithm is specified in the main function. To make it easier to play with it, it is better to do all computations in some other function, and only handle input/output in main. In your case:

solve :: Int -> Int
solve n = length $ filter chainTo89 $ [1..n - 1]

main = print (solve 1000000)

This allows us, for example, to call solve on various numbers in ghci, in tests, or in benchmarks.

benchmarks

Since you asked about algorithmic improvements, benchmarks are actually a good idea. To use the criterion benchmarking library, we add an import and replace the main function:

import Criterion.Main

-- code to benchmark

main = defaultMain
  [ bgroup "chain89"
    [ bench ("solve " ++ show n) $ whnf solve n
    | n <- [100000, 200000 .. 1000000]
    ]
  ]

This benchmarks the solve function on values below 1 million, in 100k steps. The question asks for 10 million, but i decided to run my benchmarks on the smaller inputs to get more rapid feedback. If you compile this to an executable (say, chain89), run it as chain89 --output chain89.html to produce a nich benchmark report in chain89.html. For your code as given, I get these numbers (in ms):

Looks more or less linear to me. Did you expect it to look linear? Maybe think about it a moment before reading on.

Linear?!

I didn't expect it to be linear, because we're duplicating the work in the recursive calls. You mentioned that caching the recursive results helped a lot, so I expected the complexity without caching to be clearly worse than linear, and to regain linear complexity by caching, as you would expect from a dynamic programming solution.

After seeing the benchmark results, I realized that for large numbers n, the result of sumSquareDigits n is always much smaller than n. To approximate the result after one step, consider this: if the input is n, it has k = log n / log 10 digits. In the worst case, these digits are all 9. Then the result is k * 9 * 9. For example, sumSquareDigits 999999 is only 486. So even if we multiply an already large input by 10, the result after the first step is not more than 81 bigger, which doesn't translate into too many additional recursive calls. There are additional recursive calls, so this is still slower than linear, but only very little, so we don't see it in the benchmark results.

Memoization for small inputs

We can memoize the result on small values, where small means that the values can occur after the first step. In other words, if we want to support inputs to n, we want to memoize the results for 1 to 81 * log n / log 10, because that's the biggest number we can see after the first step, see above.

Haskell supports an unusual approach to memoization: We can set up a lazy data structure which maps inputs to outputs, and in the initialization of each entry in the data structure, we access other entries. As long as we're never (directly or indirectly) accessing an entry from itself, the entries will be computed in an appropriate order.

Which data structure do we want here? We're going to use the consecutive numbers from 1 to (81 * log n / log 10) as inputs, so we probably want to use an array. We need to add an import, change chainTo89 to add the memoization and adapt solve to the new variant of chainTo89.

import Data.Array.IArray

-- other code here

makeChainTo89 :: Int -> (Int -> Bool)
makeChainTo89 n = compute where
  size = 81 * ceiling (log (fromIntegral n) / log 10)

  cache :: Array Int Bool
  cache = array (1, size) [(i, compute i) | i <- [1 .. size]]

  fetch x = cache ! x

  compute 89 = True
  compute 1 = False
  compute x = fetch (sumSquareDigits x)

solve :: Int -> Int
solve n = length $ filter chainTo89 $ [1 .. n - 1] where
  chainTo89 = makeChainTo89 n

I renamed chainTo89 to makeChainTo89 because makeChainTo89 n returns a function that computes the same values as the old chainTo89 when called with values that are less than or equal to n. The variable size holds the size of the memoization table and is computed as explained above. The cache is the memoization table itself. It contains an entry for all i between 1 and size, mapping i to compute i. The fetch function retrieves an entry from the memoization table, and the compute function does the actual work of computing the next element in the chain. Note that compute looks almost like the old chainTo89 except that it calls fetch instead of the calling itself recursively. But fetch will force a cache entry which contains a call to compute in its thunk, so that the recursion structure is actually the same. Just when the same cache entry is forced again, the result is already there.

This only works if n is big enough (so we never move out of the table after moving into it) and if there are really no other cycles than the cycles involving 1 and 89. Here's how the benchmark results look like for this version of the code:

That's almost 60% faster.

Avoiding right-associative (++)

Your digits function appends an element to the right of an immutable linked list, which requires copying the whole list. Since the order of the digits is irrelevant for the reset of the algorithm, it would be better to prepend the found digit to the list. (Even if the order is relevant, it is better to first produce the reversed list of digits, and then reverse once).

This is the new digits function:

digits :: Int -> [Int]
digits 0 = []
digits x = x `mod` 10 : digits (x `div` 10)

It is another 35% faster:

More ideas

Some other low-level ideas for speeding this up further:

• explore divMod or quotRem as an alternative to div and mod.
• explore alternatives to (^ 2).

And one high-level idea:

• avoid looping over all permutations of the same selection of digits.

Further improvements

Also check out Max Taldykin's answer for (impressive) additional improvements on top of what I propose in my answer.

Answer 2

sumSquareDigits would be better expressed as

sumSquareDigits :: Int -> Int
sumSquareDigits x = sum $ map (^2) $ digits x

Note the use of $ instead of nested parentheses to form a chain. I'm also not fond of your squareDigits, since it doesn't strictly require its input to be digits, nor does it transform its input into digits.

---

One simple change to use divMod speeds it up by about 25% (using ghc -O2):

digits :: Int -> [Int]
digits 0 = []
digits n = digits quotient ++ [remainder]
  where
    (quotient, remainder) = n `divMod` 10

Avoiding ++ saves another 8% (though it reverses the output of digits):

digits x = remainder : digits quotient

---

To get another big performance gain, you should use memoization. Memoizing just the results for 1 to 100 yields another 30% speedup (though memoizing the first 1000 results is actually worse).

digits :: Int -> [Int]
digits 0 = []
digits n = remainder : digits quotient
  where
    (quotient, remainder) = n `divMod` 10

sumSquareDigits :: Int -> Int
sumSquareDigits x = sum $ map (^2) $ digits x

chainResult :: Int -> Int
chainResult n
  | n < 100   = memo !! n
  | otherwise = chainResult $ sumSquareDigits n
  where
    chainResult' 0 = 0
    chainResult' 1 = 1
    chainResult' 89 = 89
    chainResult' n = chainResult' $ sumSquareDigits n
    memo = map (chainResult' . sumSquareDigits) [0..99]

main :: IO()
main = print $ length $ filter (\n -> 89 == chainResult n) $ [1..9999999]

The program above, which includes all of these suggestions, runs in less than half the time of your original.

Answer 3

I tried to elaborate on great answer by Toxaris. After precomputing results for first 81 * log10 n numbers, most of time is spent on extracting digits and squaring them. It is possible to calculate squares of digits for sequential numbers in incremental manner.

If you have list of squared digits for some number n then here is how to calculate list of squared digits for n+1.

inc :: [Int] -> [Int]
inc (81:xs) = 0  : inc xs
inc (x:xs)  = fromJust (lookup x nextSq) : xs
    where
        sq = [0,1,4,9,16,25,36,49,64,81]
        nextSq = zip sq $ tail sq
inc []      = [1]

To get sequence of all numbers you can iterate this function:

> take 20 $ iterate inc [1]
> [[1],[4],[9],[16],[25],[36],[49],[64],[81],[0,1],[1,1],[4,1],[9,1],[16,1],[25,1],[36,1],[49,1],[64,1],[81,1],[0,4]]

Now you just sum squares of digits and get index to precomputed cache of chain results.

Here is the full code which on my machine is 15 times faster than the original one:

import qualified Data.Vector.Unboxed as Vector
import Data.List (find, lookup)
import Data.Maybe (fromJust)
import Criterion.Main


main: IO ()
main = defaultMain [ bench "92" $ whnf p92 $ 10^7 ]


p92 :: Int -> Int
p92 n = length
    $ filter (==89)
    $ map ((cache n Vector.!) . sum)
    $ take n (iterate inc [1])


next :: Int -> Int
next = sum . map (^2) . digits


digits :: Int -> [Int]
digits 0 = []
digits x = let (a,b) = divMod x 10 in b : digits a


cache :: Int -> Vector.Vector Int
cache n
    = Vector.fromList $ 0 : map
        (fromJust . find (\x -> x == 1 || x == 89) . iterate next)
        [1 .. 81 * ceiling (logBase 10 $ fromIntegral n)]


inc :: [Int] -> [Int]
inc []      = [1]
inc (81:xs) = 0  : inc xs
inc (x:xs)  = fromJust (lookup x nextSq) : xs
    where
        sq = [0,1,4,9,16,25,36,49,64,81]
        nextSq = zip sq $ tail sq