Meta Collatz Sequence

Question

Context

Chapter 6 of Learn You A Haskell contains a function called chain, which outputs the Collatz sequence of a given input.

(In short, it takes a natural number. If that number is even, it divides it by two. If it's odd, it multiplies it by 3 and then adds 1 to that. It then takes the resulting number and applies the same thing to it, which produces a new number and so on).

Here's the book's implementation of that function:

chain :: (Integral a) => a -> [a]
chain 1 = [1]
chain n
  | even n = n : chain (n `div` 2)
  | odd n = n : chain (n * 3 + 1)

There's a function a little later in the chapter, which returns the number of chains with a starting number between 1 and 100 where the chain is longer than 15, here's that function:

numLongChains' :: Int -> Int
numLongChains' n = length (filter (\xs -> length xs > 15) (map chain [1..n]))

My code

I attempted to write a function that would output how many chains you'd have to generate in order to get the number of "long chains" equal to a certain amount. (For example, generating 100 chains produces 66 "long chains", you can see this by running numLongChains' 100).

Here's that function (written for Iron Maiden's "human number", 666):

humanNumLongChains :: (Integral a) => a -> [[a]] -> a
humanNumLongChains n xs
  | length (filter (\ys -> length ys > 15) xs) == 666 = n
  | otherwise = humanNumLongChains (n + 1) (xs ++ [chain (n + 1)])

I'm fairly new to Haskell (and purely functional programming in general), this function seems a bit messy to me, all feedback is welcome.

Answer 1

Preliminaries

Because code reviews are all about how you can improve your code, it can sometimes feel like reviewers are tearing apart the code you spent so long writing. I like to begin my review with a few words of encouragement. First, your code works and appears bug-free. Second, you seem to get Haskell enough to explore your own problems successfully. Both are no small feat -- keep it up!

Small steps

First, it's unclear to me how humanNumLongChains is supposed to be invoked, or what the function name means. From what I gather, you're looking to find the inverse of numLongChains' at 666 (i.e. what n you need to give to numLongChains' to produce 666). The n and xs are accumulator arguments which are to be initialized to specific values. It's also not immediately clear that n should be initialized to 0; I was putting 1 at first until I realized that it wouldn't check the chain for 1 if I did that.

Let's do specify the initial arguments and give our "function" a new name.

invNumLongChains1 :: Int
invNumLongChains1 = humanNumLongChains 0 []

I wrote "function" since now it's been revealed that what we really have is just a value, which appears to be 730.

I see no reason why it couldn't be a function proper. Instead of hard-coding 666, we could take that as an argument.

invNumLongChains2 :: Int -> Int
invNumLongChains2 numChains = go 0 []
  where
    go :: Int -> [[Int]] -> Int
    go n chains
      | length (filter (\ys -> length ys > 15) chains) == numChains = n
      | otherwise = humanNumLongChains (n + 1) (chains ++ [chain (n + 1)])

Note that go is a common convention for helper functions scoped locally to the main function. I'm not a huge fan of it, but I didn't want to spend longer thinking of a better name. This helper function is pretty much just humanNumLongChains copy/pasted anyway.

Now we can invoke invNumLongChains2 666 and get 730.

Efficiency

The biggest spot for inefficiency in your program is your use of lists. Unfortunately, lists in Haskell don't know their own length, so each invocation of length is O(n)!

The first simple change we can make is to instead of accumulating a list of chains, accumulate a list of chain lengths.

invNumLongChains3 :: Int -> Int
invNumLongChains3 numChains = go 0 []
  where
    go :: Int -> [Int] -> Int
    go n chainLengths
      | length (filter (\chainLength -> chainLength > 15) chainLengths) == numChains = n
      | otherwise = go (n + 1) (chainLengths ++ [length $ chain (n + 1)])

You can look at the type signature of go to quickly see the change: we've gone from taking a [[Int]] to just a regular [Int].

Unfortunately, there are two other places where lists are used inefficiently. The first is the fact that we're still using length on our filtered list and the second is that you're using (++) to append to the list. Unfortunately, (++) is also O(n) even though you're only appending a single element. This is because it needs to traverse the whole linked list to append it.

It's not difficult to see that we don't even need the list of lengths. We can just keep track of a single number representing how many long chains we've seen so far.

invNumLongChains4 :: Int -> Int
invNumLongChains4 numChains = go 0 0
  where
    go :: Int -> Int -> Int
    go n numLongChains
      | numLongChains == numChains = n
      | otherwise = if length (chain (n + 1)) > 15
                      then go (n + 1) (numLongChains + 1)
                      else go (n + 1) numLongChains

Now go has gone from taking a [Int] to a single Int, and our argument has changed from [] to 0 to reflect this.

Getting lazier

I'm running out of steam a little, but I wanted to leave two alternate versions of your function that make use of laziness. I think they are both equivalently efficient to invNumLongChains4 (maybe a little less, depending on how the compiler optimizes things), but they turn the problem on its head a little bit.

Instead of iterating over n and accumulating a count of long chain lengths, the idea is to create an infinite list where the element at position n represents the number of long chains at n. Then we just find the index of 666 in that infinite list. This works because Haskell is lazy, so we only evaluate the elements of this list when we check them!

Some notes: Because elemIndex returns a Maybe, we need to extract the value using fromJust. This is usually bad practice (you would want to instead keep the result a Maybe) because it will error if there is no element. However, because we're using it on an infinite list we will never get a Nothing back (the code will just infinitely loop if it can't find a match). The other part is that elemIndex 0-indexes and we're starting our list from 1, so we need to add 1 to the result it finds.

Here is the first rewrite. I tried to break things apart enough to make it clear what's going on.

invNumLongChains5 :: Int -> Int
invNumLongChains5 numChains = (+1) $ fromJust $ elemIndex numChains longChainCounts
  where
    chains :: [[Int]]
    chains = map chain [1..]
    countLongChains :: Int -> [[Int]] -> [Int]
    countLongChains prevCount (chain:chains)
      | length chain > 15 = prevCount + 1 : countLongChains (prevCount + 1) chains
      | otherwise = prevCount : countLongChains prevCount chains
    longChainCounts :: [Int]
    longChainCounts = countLongChains 0 chains

In the second rewrite, I make use of a standard function called scanl1 to replace the helper function countLongChains. It's all right if this doesn't make sense to you! If you're really interested in understanding how it works, let me know in a comment and I can try to pull together some time to explain.

invNumLongChains6 :: Int -> Int
invNumLongChains6 numChains = (+1) . fromJust $ elemIndex numChains longChainCounts
   where
     isLongEnough :: [Int] -> Int
     isLongEnough chain
       | length chain > 15 = 1
       | otherwise = 0
     longChainCounts :: [Int]
     longChainCounts = scanl1 (+) $ map (isLongEnough . chain) [1..]

I'm no Haskell guru, but I would probably prefer invNumLongChains4 to invNumLongChains6 for readability's sake. However, I do think that the latter could be considered a more "functional" approach, which is why I included it.

Try it Online!

Here's a TIO snippet with all the functions, in case you wanted to play around with them. Or, more likely, in case I forgot to copy/paste an update I made to my answer -- the TIO snippet should have everything up-to-date.

Try it Online!

Postscripts

• I use Int instead of (Integral a) => a because there are places where I compare the variable with that type to length x. For historical reasons, length x returns an Int. It's a little annoying to have to put fromIntegral so I instead just use narrow the type to Int.
• I would probably also want to pull out the magic number 15 that is present for the length, but decided against it for simplicity's sake.
• There's another major point on efficiency: memoizing chain. But I decided that this was too much to discuss in my review.
• I would recommend as an exercise that you test the efficiency of all of these rewrites. I didn't, and I'm somewhat-maybe-slightly worried I'm wrong about my efficiency assessments (especially of invNumLongChains5 and invNumLongChains6).