Fast sums of power algorithm

Question

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

Completely changing the code is not fair.. I've rolled back Rev 3 → 2. — 200_success
– 200_success, Commented Feb 22, 2015 at 16:59
Links can rot. Please include a description of the challenge here in your question. — 200_success
– 200_success, Commented Feb 22, 2015 at 17:05

Joachim Breitner · Accepted Answer · 2015-02-22 11:24:20Z

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!

Dulguun Otgon · Accepted Answer · 2015-02-23 03:54:33Z

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

Stack Exchange Network

Fast sums of power algorithm

2 Answers 2

You must log in to answer this question.

Hot Network Questions

Fast sums of power algorithm

2 Answers 2

You must log in to answer this question.

Related

Hot Network Questions