programming challenge – HackerRank – Binary tree to compute the Nth power of natural numbers that sum up to X


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.