First of all, we can make your prime_factors_mult
easier to read if we use where
:
prime_factors_mult :: Int -> [(Int, Int)]
prime_factors_mult n = reverse . fst . foldl go ([], n) $ candidates
where
candidates = [2 .. floor . sqrt . fromIntegral $ n]
go acc@(factors, val) x
| val `mod` x == 0 = let (k, v) = reduce val x in ((x, k) : factors, val `div` v)
| otherwise = acc
By the way, Haskell programs use camelCase
, so primFactorsMult
would be appropriate. However, does this function really capture your intend? Let us write it as imperative variant for a second:
# pseudo-python
def prime_factors_mult(n):
result = []
for i in range(2, sqrt(n)):
count = 0
while n % i == 0:
n = n / i
count = count + 1
if count > 0:
result = result.append({i, count})
return result
This is a lot easier to read. We should strive to have an algorithm in Haskell that's just as easy to read, right? By the way, both the pseudo-code as well as your original code contained an error: what happens on 2*7
? \$\sqrt{14} < 7\$, therefore we never check the second prime. We should check that too:
# pseudo-python
def prime_factors_mult(n):
result = []
for i in range(2, sqrt(n)):
count = 0
while n % i == 0:
n = n / i
count = count + 1
if count > 0:
result = result.append({i, count})
# n is prime at this point:
if n > 1:
result = result.append({n, 1})
return result
Allright. Now let's try to get a variant of Haskell that's just as readable. We do that by explicit recursion and then later figure out how we can get rid of that:
isDivisibleBy n k = n `rem` k == 0
primeFactorsMult :: Int -> [(Int, Int)]
primeFactorsMult = go 2
where
go _ 1 = []
go k n
| n `isDivisibleBy ` k = (k, m) : go (k + 1) (n `div` d)
| n < k * k = [(n, 1)]
| otherwise = go (k + 1) n
where
(m, d) = reduce n k
Allright, that seems to work. However, do we actually need the isDivisibleBy
check? No. reduce
does that for us. So let's get rid of that:
primeFactorsMult :: Int -> [(Int, Int)]
primeFactorsMult = filter ((0 /=) . snd) . go 2
where
go _ 1 = []
go k n
| n <= k = [(n, 1)]
| otherwise = (k, m) : go (k + 1) (n `div` d)
where
(m, n') = reduce n k
Great. Now we're getting somewhere. We have now two paths we can choose to continue. Either we take our new approach, which is essentially mapAccumL
, or we inline reduce
. Let us follow both.
Mapping with a state
In the code above, we essentially mapped over the list of factors [2..n]
with a state: the current value of n
, or rather n `div` d
. It turns out that there is already a function for that kind of operation, namely mapAccumL
. We can incorporate into our function:
primeFactorsMult :: Int -> [(Int, Int)]
primeFactorsMult n = filter ((0 /=) . snd) . snd . mapAccumL go n $ 2 : [3,5..n]
where
go n k = (n `div` d, (k, m))
where
(m, n') = reduce n k
As you can see, we got rid of the recursive calls in go
and the base case, since that's now handled by mapAccumL
. We also swapped the arguments, since mapAccumL
gives the state first. Also, since primes are odd (except for 2
), we've optimized the list slightly.
Unfortunately, we still need to get rid of those with 0
multiplicity. However, what if we didn't return a pair?
Return factors, one by one
This time, we inline the functionality of reduce
:
primeFactors :: Int -> [Int]
primeFactors n = go 2
where
go _ 1 = []
go k n = case n `quotRem` k of
(q, 0) -> k : go k q
_ -> go (k + 1) n
This will now return a list like [2,2,2,2,2,2,2,2,2,5,5,5,7,7,9,9,10]
. By the way, this quotRem
shows a possible optimization of reduce
:
reduce :: Int -> Int -> (Int, Int)
reduce val x =
if m == 0 then
let (k, v) = reduce d x in (1 + k, v * x)
else
(0, 1)
where
(d, m) = val `divMod` x
But when all your numbers are positive, you want to use quotRem
since it's usually faster. Note that many modern processors do mod
and div
at the same time, so using the appropriate function when you ened both results can be a boon.
Back to our prime factors. Now we have the factors. But how can we get the multiplicity? We use group
:
primeFactorsMult :: Int -> [(Int, Int)]
primeFactorsMult = map (\xs -> (head xs, length xs)) . group . primeFactors
Since group
returns only non-empty lists, we can safely apply head
.
Remarks on style
I strongly suggest you to move large lambdas into their own bindings. That made the first variant of prime_factors_mult
more readable and is the usual style, for example in GHC's code.
Also, if you use both the results of div
and mod
use divMod
, or—if you use positive numbers only or don't want the specific behaviour behind divMod
—quotRem
.
Don't be afraid to give something in your code a name. That makes it easier to grasp and to change later. Also, you might notice that a function is reusable.
And, by the way, primes are a "hard" topic, as in: it's not easy to get proper bounds of the numbers you have to check, since you would need to know whether n
is already a prime, or whether n / 2
is a prime, or n / 3
and so on.
Exercises
- In
go
you can use (k + 2)
if you handle 2
separately. How can you do that?
- In most
go
functions we can stop much, much sooner if n
is a prime if we know the original n
. Why? (Hint: what is the largest prime factor that can be in a number? And what is the second largest prime factor?)