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 , 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):

• multigram
• relatives
• magical3
• perfectpowers
• factovisors

In a subsequent post I’ll continue on the number theory theme and talk about modular arithmetic.