Number theory is a topic that comes up fairly regularly in competitive programming, and it’s a very nice fit for Haskell. I’ve developed a bunch of code over the years that regularly comes in handy. None of this is particularly optimized, and it’s definitely no match for a specialized library like arithmoi, but in a competitive programming context it usually does the trick!
A few imports first:
import Control.Arrow
import Data.List (group, sort)
import Data.Map (Map)
import qualified Data.Map as M
Primes
We start with a basic definition of the list of primes, made with a simple recursive sieve, but with one very big optimization: when we find a prime , instead of simply filtering out all the multiples of
in the rest of the list, we first take all the numbers less than
and pass them through without testing; composite numbers less than
would have already been filtered out by a smaller prime.
primes :: [Integer]
primes = 2 : sieve primes [3..]
where
sieve (p:ps) xs =
let (h,t) = span (< p*p) xs
in h ++ sieve ps (filter ((/=0).(`mod`p)) t)
I got this code from the Haskell wiki page on prime numbers. On my machine this allows us to find all the primes up to one million in about 4 seconds. Not blazing fast by any means, and of course this is not actually a true sieve—but it’s short, relatively easy to remember, and works just fine for many purposes. (There are some competitive programming problems requiring a true sieve, but I typically solve those in Java. Maybe someday I will figure out a concise way to solve them in Haskell.)
Factoring
Now that we have our list of primes, we can write a function to find prime factorizations:
listFactors :: Integer -> [Integer]
listFactors = go primes
where
go _ 1 = []
go (p:ps) n
| p*p > n = [n]
| n `mod` p == 0 = p : go (p:ps) (n `div` p)
| otherwise = go ps n
This is relatively straightforward. Note how we stop when the next prime is greater than the square root of the number being tested, because if there were a prime factor we would have already found it by that point.
…and related functions
Finally we can use listFactors
to build a few other useful functions:
factor :: Integer -> Map Integer Int
factor = listFactors >>> group >>> map (head &&& length) >>> M.fromList
divisors :: Integer -> [Integer]
divisors = factor >>> M.assocs >>> map (\(p,k) -> take (k+1) (iterate (*p) 1))
>>> sequence >>> map product
totient :: Integer -> Integer
totient = factor >>> M.assocs >>> map (\(p,k) -> p^(k-1) * (p-1)) >>> product
factor
yields a Map
whose keys are unique primes and whose values are the corresponding powers; for example, factor 600 = M.fromList [(2,3), (3,1), (5,2)]
, corresponding to the factorization . It works by grouping together like prime factors (note that
listFactors
guarantees to generate a sorted list of prime factors), counting each group, and building a Map
.
divisors n
generates a list of all divisors of n
. It works by generating all powers of each prime from up to
, and combining them in all possible ways using
sequence
. Note it does not guarantee to generate the divisors in order.
totient
implements the Euler totient function: totient n
says how many numbers from 1
to n
are relatively prime to n
. To understand how it works, see this series of four blog posts I wrote on my other blog: part 1, part 2, part 3, part 4.
Problems
Here are a few problems for you to try (ordered roughly from easier to more difficult):
In a subsequent post I’ll continue on the number theory theme and talk about modular arithmetic.
Pingback: Resumen de lecturas compartidas del 1 al 8 de febrero de 2020 | Vestigium
Thanks for your post, Brent!
Can you spend a few words on how “primes” works with its two points of recursion?
Sure! It is somewhat mindbending to be sure. I’ll give some thought to explaining it well and turn it into another blog post.
Excellent! :-)
Looking forward to it.
Pingback: Competitive programming in Haskell: summer series | blog :: Brent -> [String]