0
\$\begingroup\$

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 = []
      | goal (head s) = head s : search' (tail 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
\$\endgroup\$
2
2
\$\begingroup\$
  • 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!

\$\endgroup\$
0
\$\begingroup\$

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
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.