# Competitive Programming in Haskell: primes and factoring

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 \(p\), instead of simply filtering out all the multiples of \(p\) in the rest of the list, we first take all the numbers less than \(p^2\) and pass them through without testing; composite numbers less than \(p^2\) 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 \(600 = 2^3 \cdot 3 \cdot 5^2\). 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 \(p^0\) up to \(p^k\), 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.