HackerRank - Binary tree to compute the Nth power of natural numbers that sum up to X

Question

I'm trying to solve this challenge on HackerRank.

In short, given x and n, I have to determine how many ways I can pick numbers from the nth powers of the natural numbers such that their sum is x.

So my first reasoning was to compute the list of the nth powers that I need as a first step. This list is

list = takeWhile (<= x)  [i^n | i <- [1..]]`

Now list is the bucket from which I have to pick numbers to add up to n.

I first thought about list comprehensions, because that's how I would pick a given number of numbers from the list; for instance this is how I would generate the sumList of the sums of any different triplets from list:

sumList = [x + y + z | x <- list, y <- dropWhile (<= x) list, z <- dropWhile (<= y) list]

Then I would just need to length $ filter (== x) sumList.

However, after some thinking, I concluded that list comprehensions are not the way to go, since I don't know in advance how many numbers I have to pick from list when attempting to sum them up to x.

So I thought that any possible sum corresponds to a combination of pick or not-pick while traversing the list. This made me think of binary trees, and I eventually came up with this solution, which fails the Test Case 3 for timeout:

main :: IO()
main = do
  [x,n] <- sequence $ replicate 2 ((read :: String -> Int) <$> getLine)
  print $ length $ filter (== x) $ getLeaves $ listToSumTree $ takeWhile (<= x) [i^n | i <- [1..]]

data Tree a = Leaf a | Node a (Tree a) (Tree a) deriving (Show)

listToTreeImpl :: (Num a, Eq a) => [a] -> Tree a -> Tree a
listToTreeImpl [] tree = tree
listToTreeImpl (l:ls) (Leaf x) = Node x
                         (listToTreeImpl ls $ Leaf x)
                         (listToTreeImpl ls $ Leaf (x + l))

getLeaves :: Tree a -> [a]
getLeaves (Leaf a) = [a]
getLeaves (Node _ (Leaf _) (Node _ _ _)) = error "this should not happen"
getLeaves (Node _ (Node _ _ _) (Leaf _)) = error "this should not happen"
getLeaves (Node _ left right) = getLeaves left ++ getLeaves right

listToSumTree :: (Num a, Eq a) => [a] -> Tree a
listToSumTree ls = listToTreeImpl ls (Leaf 0)

I can only think of getLeaves $ listToSumTree as the only critical part of the program, since listToSumTree constructs a full tree, which is not needed (only the leaves are), and getLeaves traverses all the tree, only to get to the bottom layer of it.

Answer 1

What you're essentially doing with your binary tree is checking all possible combinations of numbers from the list [i^n | i <- [1..x]]. That's quite a lot of combinations! (Exercise: how many?)

Not all of those combinations are sensible, however. For example, if x = 100 and n = 2, and you've already chosen, say, 4 and 5 (giving 6²+7²=85), then adding 8²=64 is not going to help, nor is adding 9²=81. In your current solution, however, you're still checking those.

So how does one encode this knowledge in the search for solutions? The key thing is that we can build the solution recursively: we can use an already computed answer for a smaller case to help answer a larger case. If we're given x = 10 and n = 2, then we can only add 1², 2² or 3². If we try adding 3², we have only 1 left to fill, which we can only do by adding 1². If, instead, we start by trying 2², we have to make 6, which is impossible. Starting with 1² lets us choose 3² to fill up the remaining 9. Since only the set of numbers counts, not their order, this means there is only one solution for x=10,n=2: namely, {1²,3²} (as stated in the problem).

Checking whether two solutions are the same modulo ordering is annoying, however. Can we do better? Yes we can, by adding another parameter to the problem.

Think about how we could do this before reading on.

---

A way to change the problem to prevent us from generating multiple solutions that are really the same set of numbers, is to add as a parameter the maximum value we're allowed to choose. So a problem is now not only specified by (x=10,n=2), but also by an integer M; for example, if M=9, then we may add 1²≤9, 2²≤9 and 3²≤9, so (x=10,n=2,M=9) has one solution: {1²,3²}. However, (x=10,n=2,M=2) has no solutions, because by only using 1² and 2² you can never make 10.

How is this useful? Well, consider how we could, as a lazy office worker, go about computing how many solutions there are to (x=10,n=2,M=10). (For the original problem, we just set M=x, since choosing numbers higher than x is certainly useless.) This lazy office worker realises that either x=0, in which case we have precisely 1 solution (choosing no numbers at all), or x≠0, in which case we need to choose at least one number. If we choose, say, i², then the remainder still to fill is 10 - i². Since our lazy office worker is also clever, they realise that if we force any solutions for 10 - i² to only choose numbers smaller than i² (i.e. set M = i² - 1), then we won't ever get any duplicates: there may certainly be solutions for our (x,n,M) that choose i² and also some larger j², but then we'll find that one by starting with j² and then going downward.

So our lazy office worker writes down a series of smaller problems: [(x-in,n,in-1) | i <- [1..M]], to solve for the next office worker in line. (These are recursive calls.) The total number of solutions is just the sum of the number of solutions for each of these subproblems.

An alternative, more direct way to explain the same way of solving, is to note that while all solutions are sets of distinct integers, we can represent those sets by sorted lists. If the empty solution is not possible (i.e. x≠0), then each of the possible solutions for our current problem has one largest value, so we can find all solutions for the problem going over each largest value iⁿ from 1 to M, computing the solutions for (x - i², n, iⁿ - 1), and prepending that largest value i² to its corresponding sub-solutions. If, instead of building the actual solution lists we just compute the number of them, we arrive at the same solution as above.

---

Now the above is not very formalised yet, and that is on purpose. Because for me, the next step to formalisation would be writing out the recursive function, which is 95% of the work in implementing the program. And that's your task. :)

If any of the above is unclear (probably), or you get stuck, feel free to ask further questions.

EDIT: When I coded out my first solution for the problem, I also cached recursive calls in a table, making it a dynamic programming solution. However, for the given input range that apparently doesn't help, so a standard recursive program works fine.

EDIT 2: If you want to handle 1000 1, then you do need dynamic programming :)