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.

About Brent

Assistant Professor of Computer Science at Hendrix College. Functional programmer, mathematician, teacher, pianist, follower of Jesus.
This entry was posted in haskell and tagged , , , , . Bookmark the permalink.

5 Responses to Competitive Programming in Haskell: primes and factoring

  1. Pingback: Resumen de lecturas compartidas del 1 al 8 de febrero de 2020 | Vestigium

  2. Matthias Braun says:

    Thanks for your post, Brent!
    Can you spend a few words on how “primes” works with its two points of recursion?

  3. Pingback: Competitive programming in Haskell: summer series | blog :: Brent -> [String]

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Google photo

You are commenting using your Google account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s

This site uses Akismet to reduce spam. Learn how your comment data is processed.