Efficiently calculating perfect powers in Haskell

Question

I'm trying to calculate perfect powers to solve this code wars challenge: https://www.codewars.com/kata/55f4e56315a375c1ed000159/haskell

The number 81 has a special property, a certain power of the sum of its digits is equal to 81 (nine squared). Eighty one (81), is the first number in having this property (not considering numbers of one digit). The next one, is 512. Let's see both cases with the details

\$8 + 1 = 9\$ and \$9^2\$ = 81

\$512 = 5 + 1 + 2 = 8\$ and \$8^3\$ = 512

We need to make a function, power_sumDigTerm(), that receives a number n and may output the n-th term of this sequence of numbers. The cases we presented above means that

power_sumDigTerm(1) == 81

power_sumDigTerm(2) == 512

The initial brute force approach where I check the digits sum for every number could only get up to the 11th term in the sequence in reasonable time.

I was able to improve on this by only checking numbers that are perfect powers

Using this approach I was able to get up to the 17th number in the sequence.

This is the code that I am using (apologies for the poor coding style, I'm a Haskell beginner with an imperative programming background):

module Codewars.G964.Powersumdig where
import qualified Data.Vector.Unboxed as V
import Data.Vector.Unboxed (Vector, fromList, (!), snoc, foldl', findIndex, accum)
import Data.Int (Int64)
import Data.Maybe (fromJust)

v1 (a,_,_) = a
v2 (_,a,_) = a
v3 (_,_,a) = a

inc_exp :: (Int64, Int, Int) -> Int -> (Int64, Int, Int)
inc_exp (_, b, c) _ = ((fromIntegral c)^nb, nb, c)
  where nb = succ b
  
-- next_pp : get the next "perfect power"
-- The vector input contains a tuple of potential next perfect powers
-- With the base and exponents used to calculate them
-- Returns a tuple of the next perfect power, the exponent used to calculate the power
-- and the updated Vector state to get the next perfect power
next_pp :: Vector (Int64, Int, Int) -> (Int64, Int, Vector (Int64, Int, Int))
next_pp v = let next  = foldl' ((. v1) . min) maxBound v
                i     = fromJust $ findIndex ((next ==) . v1) v -- fromJust should be safe as vector should contain it
                upd_v = accum inc_exp v [(i,0)]
                last  = upd_v ! ((V.length upd_v)-1) -- (last upd_v) doesn't work??
                new_v = if v2 last == 3 then upd_v `snoc` ((fromIntegral (succ (v3 last)))^2, 2, succ (v3 last)) else upd_v
            in (next, v2 (v ! i), new_v)

get_next :: Int64 -> Int -> Vector (Int64, Int, Int) -> Int64
get_next n d v
  | is_valid && d == 1 = np
  | is_valid = get_next np (d-1) nps
  | otherwise = get_next np d nps
    where (np, e, nps) = next_pp v
          is_valid     = (sum_digits np)^e == np
          
sumd 0 acc = acc
sumd x acc = sumd d (acc + m)
  where (d,m) = x `divMod` 10

sum_digits x = sumd x 0

powerSumDigTerm :: Int -> Integer
-- 16 is the first perfect power with 2 or more digits
-- the vector is initialised with values to generate the next perfect power (which happens to also be 16)
powerSumDigTerm n = toInteger $ get_next 16 n (fromList [(maxBound,0,0),(maxBound,0,1),(32,5,2),(27,3,3),(16,2,4)])

I used ghc profiling tools to find which parts of the code are the bottlenecks.

From this I can tell that:

• 99.8% of the time is spent in the next_pp function.
• 39.6% of time is spent calculating i
• 39.2% of time is spent calculating next
• 12.4% of time is spent calculating upd_v
• 8.3% of time is spent calculating new_v

I need to be able to calculate up to the 40th code in the sequence, so this code is woefully inadequate, but I'm stuck for where to go next with it. Any suggestions? Code style suggestions are also welcome.

Perhaps there is a better set of numbers to check that is more specific than perfect powers?

I also tried using sequences and boxed vectors but unboxed vectors performed best out of the ones I tried. (I am worried that later terms in the sequence could overflow a 64 bit int)

I tried using sequences unstableSort to sort the list each time and find the smallest power but this ended up being slower than using unboxed vectors and folds to find the smallest power.

Edit: my second attempt using Data.Heap (unfortunately codewars doesn't support Data.Heap though)

module Codewars.G964.Powersumdig where
import Data.Heap
import Data.Maybe (fromJust)

inc_exp :: (Int, Int) -> (Integer, (Int, Int))
inc_exp (e, b) = ((fromIntegral b)^ne, (ne, b))
  where ne = succ e
  
next_pp :: MinPrioHeap Integer (Int, Int) -> (Integer, Int, MinPrioHeap Integer (Int, Int))
next_pp v = let next = fromJust  $ viewHead h
                tmp_h = fromJust $ viewTail h
                base = snd $ snd next
                exp = fst $ snd next
                upd_h = insert (inc_exp $ snd next) tmp_h
                new_h = if exp == 2 then insert ((fromIntegral $ succ base)^2, (2, succ base)) upd_h else upd_h
             in (fst next, base, new_h)

get_next :: Integer -> Int -> MinPrioHeap Integer (Int, Int) -> Integer
get_next n d h
  | n < 10 = get_next np d nps
  | is_valid && d == 1 = np
  | is_valid = get_next np (d-1) nps
  | otherwise = get_next np d nps
    where (np, b, nps) = next_pp h
          is_valid     = fromIntegral (sum_digits np) == b
          
sumd 0 acc = acc
sumd x acc = sumd d (acc + m)
  where (d,m) = x `divMod` 10

sum_digits x = sumd x 0

powerSumDigTerm :: Int -> Integer
powerSumDigTerm n = toInteger $ get_next 1 n (fromList [(4, (2,2))])

Answer 1

So, it turns out that there really isn't a nice way to solve this in the general case as far as I can tell.

If you know an upper bound for the values of the base (b) and the exponent (e) then you can generate an (unsorted) list of all perfect powers, which only takes O(b*e) time.

Then you only have to sort the list once and filter the list once.

import Data.List
          
sumd 0 acc = acc
sumd x acc = sumd d (acc + m)
  where (d,m) = x `divMod` 10

sum_digits x = sumd x 0

perfect_powers :: Int -> Int -> [(Integer,Int)]
perfect_powers b e = go 2 2 []
  where
    go x y l
      | x == b && y == e = l
      | y == e = go (x+1) 2 (((toInteger x)^y,y):l)
      | otherwise = go x (y+1) (((toInteger x)^y,y):l)

powerSumDigTerm :: Int -> Integer
powerSumDigTerm n = fst $ (filter (\(x,y) -> (sum_digits x)^y == x) $ sortOn fst $ perfect_powers 200 100) !! (n-1)

Overall, I'm kind of unsatisfied with this solution, if anyone knows one that works in the general case (i.e without bounds on base and exponent) please let me know