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
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.)
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
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
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.
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.