# Fast sums of power algorithm

This code reaches time limit on 1 test case, solving this challenge. The approach that I'm using is from the book Algorithms Functional Programming Approach, which is a backtracking depth searching method. How do I improve speed of this code? Or is there better way to do it?

searchDfs :: (Eq node) =>
(node -> [node]) -> -- generates successor nodes
(node -> Bool) -> -- is goal
node -> -- first node
[node]
searchDfs succ goal x = search' ([x])
where
search' s
| null s = []
| otherwise = let x = head s
in search' (succ x ++ tail s)

type Degree = Int
type CurrentSum = Int
type CurrentNumber = Int
type Node = (CurrentSum,CurrentNumber,Degree,Goal)
type Goal = Int

succN :: Node -> [Node]
succN (sum,i,d, goal)  = [( sum+i'^d, i', d, goal) | i' <- [(i+1)..goal], sum+i'^d <= goal ]

goal :: Node -> Bool
goal (sum,_,_,m) = sum == m

countSols :: Degree -> Goal -> Int
countSols d m = foldr (+) 0 $map (\(sum,i) -> length (searchDfs succN goal (sum,i,d,m))) startingNumbers where startingNumbers = [ (i^d,i) | i <- [1..m], i^d <= m] main = do m <- readLn d <- readLn let c = countSols d m print c  ## 2 Answers • The search' could could be written more idiomatically (but not necessarily faster) as search' [] = [] search' (x:xs) | goal x = x : search' xs | otherwise = search' (succ x ++ xs)  • The Eq constraint on searchDfs seems to be unused. • In succN, you are calculating sum+i'^d twice (unless the compiler optimizes that away, but it’s better not to rely on that). Also, it will calculate it for all values of i' until goal, although you probably want to abort early. Maybe using takeWhile would be better:  succN (sum,i,d, goal) = takeWhile (\(s,_,_,_) -> s <= goal) [( sum+i'^d, i', d, goal) | i' <- [(i+1)..goal]]  • It looks as if the d and goal field of your Nodes is always the same. Then you shouldn’t have it as part of the Node, but rather pass it around explicitly. You can do that without modifying searchDfs! It was because of i^d generating too big number that Int can't hold, it became unresponsive/froze. Also because computing same i^d on different nodes was making it slower. Now it computes them once. Then uses map lookup. import Control.Monad (guard) import qualified Data.Map.Strict as Map import Data.Map ((!),fromList,Map(..),member) -- depth first search higher order function depthFirstSearch :: (node -> [node]) -> -- successor node generator (node -> Bool) -> -- is goal checker [node] -> -- starting root nodes [node] -- goal nodes depthFirstSearch succ goal roots = search' roots where search' [] = [] search' (x:xs) | goal x = x : search' xs | otherwise = search' (succ x ++ xs) type Degree = Int type CurrentSum = Int type CurrentNumber = Int type Node = (CurrentNumber,CurrentSum) type Goal = Int -- generates valid successor nodes succN :: Goal -> Map Int Int -> Node -> [Node] succN goal _map (i,sum) = do i' <- [(i+1)..goal] guard (member i' _map) let sum' = sum + _map!i' guard (sum' <= goal) return (i',sum') -- checks if the node is the goal goalN :: Goal -> Node -> Bool goalN goal (_,sum) = goal == sum -- counts solutions solCount :: Degree -> Goal -> Int solCount d goal = let roots = rts d goal -- [(i,i^d) | i <- [1..goal], i^d <= goal] _map = Map.fromList roots nodes = depthFirstSearch (succN goal _map) (goalN goal) roots c = length nodes in c rts :: Degree -> Goal -> [Node] rts d' g' = do let d = fromIntegral d' g = fromIntegral g' a <- [(1::Integer)..(fromIntegral g)] let aNth = a^d guard (aNth <= g) let bNth = (fromInteger :: Integer -> Int) aNth b = (fromInteger :: Integer -> Int) a return (b,bNth) main = do g <- readLn d <- readLn print$ solCount d g