Competitive programming in Haskell: Enumeration

I’m in the middle of a multi-part series on implementing BFS in Haskell. In my last post, we saw one implementation, but I claimed that it is not fast enough to solve Modulo Solitaire, and I promised to show off a faster implementation in this post, but I lied; we have to take a small detour first.

The main idea to make a faster BFS implementation is to replace the HashMaps from last time with mutable arrays, but hopefully in such a way that we get to keep the same pure API. Using mutable arrays introduces a few wrinkles, though.

1. The API we have says we get to use any type v for our vertices, as long as it is an instance of Ord and Hashable. However, this is not going to work so well for mutable arrays. We still want the external API to allow us to use any type for our vertices, but we will need a way to convert vertices to and from Int values we can use to index the internal mutable array.

2. A data structre like HashMap is dynamically sized, but we don’t have this luxury with arrays. We will have to know the size of the array up front.

In other words, we need to provide a way to bijectively map vertices to a finite prefix of the natural numbers; that is, we need what I call invertible enumerations. This idea has come up for me multiple times: in 2016, I wrote about using such an abstraction to solve another competitive programming problem, and in 2019 I published a library for working with invertible enumerations. I’ve now put together a lightweight version of that library for use in competitive programming. I’ll walk through the code below, and you can also find the source code in my comprog-hs repository.

First, some extensions and imports.

{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module Enumeration where

import qualified Data.List as L
import Data.Hashable
import qualified Data.Array as A
import qualified Data.HashMap.Strict as HM

An Enumeration a consists of a cardinality, and two functions, select and locate, which together form a bijection between (some subset of) values of type a and a finite prefix of the natural numbers. We can convert an Enumeration into a list just by mapping the select function over that finite prefix.

data Enumeration a = Enumeration
  { card   :: !Int
  , select :: Int -> a
  , locate :: a -> Int
  }

enumerate :: Enumeration a -> [a]
enumerate e = map (select e) [0 .. card e-1]

Since a occurs both positively and negatively, Enumeration is not a Functor, but we can map over Enumerations as long as we provide both directions of a bijection a <-> b.

mapE :: (a -> b) -> (b -> a) -> Enumeration a -> Enumeration b
mapE f g (Enumeration c s l) = Enumeration c (f . s) (l . g)

We have various fundamental ways to build enumerations: empty and unit enumerations, and an identity enumeration on a finite prefix of natural numbers.

voidE :: Enumeration a
voidE = Enumeration 0 (error "select void") (error "locate void")

unitE :: Enumeration ()
unitE = singletonE ()

singletonE :: a -> Enumeration a
singletonE a = Enumeration 1 (const a) (const 0)

finiteE :: Int -> Enumeration Int
finiteE n = Enumeration n id id

We can automatically enumerate all the values of a Bounded Enum instance. This is useful, for example, when we have made a custom enumeration type.

boundedEnum :: forall a. (Enum a, Bounded a) => Enumeration a
boundedEnum = Enumeration
  { card   = hi - lo + 1
  , select = toEnum . (+lo)
  , locate = subtract lo . fromEnum
  }
  where
    lo, hi :: Int
    lo = fromIntegral (fromEnum (minBound @a))
    hi = fromIntegral (fromEnum (maxBound @a))

We can also build an enumeration from an explicit list. We want to make sure this is efficient, since it is easy to imagine using this e.g. on a very large list of vertex values given as part of the input of a problem. So we build an array and a hashmap to allow fast lookups in both directions.

listE :: forall a. (Hashable a, Eq a) => [a] -> Enumeration a
listE as = Enumeration n (toA A.!) (fromA HM.!)
  where
    n = length as
    toA :: A.Array Int a
    toA = A.listArray (0,n-1) as
    fromA :: HM.HashMap a Int
    fromA = HM.fromList (zip as [0 :: Int ..])

Finally, we have a couple ways to combine enumerations into more complex ones, via sum and product.

(>+<) :: Enumeration a -> Enumeration b -> Enumeration (Either a b)
a >+< b = Enumeration
  { card   = card a + card b
  , select = \k -> if k < card a then Left (select a k) else Right (select b (k - card a))
  , locate = either (locate a) ((+card a) . locate b)
  }

(>*<) :: Enumeration a -> Enumeration b -> Enumeration (a,b)
a >*< b = Enumeration
  { card = card a * card b
  , select = \k -> let (i,j) = k `divMod` card b in (select a i, select b j)
  , locate = \(x,y) -> card b * locate a x + locate b y
  }

There are a few more combinators in the source code but I don’t know whether I’ll ever use them. You can read about them if you want. For now, let’s try using this to solve a problem!

…ah, who am I kidding, I can’t find any problems that can be directly solved using this framework. Invertibility is a double-edged sword—we absolutely need it for creating an efficient BFS with arbitrary vertices, and the combinators will come in quite handy if we want to use some complex type for vertices. However, requiring invertibility also limits the expressiveness of the library. For example, there is no Monad instance. This is why my simple-enumeration library has both invertible and non-invertible variants.