## Lazily Grouping in Haskell

Here’s a cool trick:

```
minimum :: Ord a => [a] -> a
minimum = head . sort
```

This is $\mathcal{O}(n)$ in Haskell, not $\mathcal{O}(n \log n)$ as you might expect. And this isn’t because Haskell is using some weird linear-time sorting algorithm; indeed, the following is $\mathcal{O}(n \log n)$:

```
maximum :: Ord a => [a] -> a
maximum = last . sort
```

No: since the implementation of `minimum`

above only demands the first element of the list, and since `sort`

has been carefully implemented, only a linear amount of work will be done to retrieve it.

It’s not easy to structure programs to have the same property as `sort`

does above: to be maximally lazy, such that unnecessary work is not performed. Today I was working on a maximally lazy implementation of the following program:

```
groupOn :: Eq k => (a -> k) -> [a] -> [(k,[a])]
= ...
groupOn
>>> groupOn (`rem` 2) [1..5]
1,[1,3,5]),(0,[2,4])]
[(
>>> groupOn (`rem` 3) [5,8,3,6,2]
2,[5,8,2]),(0,[3,6])] [(
```

This function groups the elements of a list according to some key function. The desired behaviour here is a little subtle: we don’t want to just group adjacent elements, for instance.

``rem` 3) [5,8,3,6,2] ≢ [(2,[5,8]),(0,[3,6]),(2,[2])] groupOn (`

And we don’t want to reorder the elements of the list by the keys:

``rem` 3) [5,8,3,6,2] ≢ [(0,[3,6]),(2,[5,8,2])] groupOn (`

These constraints make it especially tricky to make this function lazy. In fact, at first glance, it seems impossible. What should, for instance, `groupOn id [1..]`

return? It can’t even fill out the first group, since it will never find another `1`

. However, it *can* fill out the first key. And, in fact, the second. And it can fill out the first element of the first group. Precisely:

`id [1..] ≡ [(1,1:⊥), (2,2:⊥), (3,3:⊥), ... groupOn `

Another example is `groupOn id (repeat 1)`

, or `groupOn id (cycle [1,2,3])`

. These each have partially-defined answers:

```
id (repeat 1) ≡ (1,repeat 1):⊥
groupOn
id (cycle [1,2,3]) ≡ (1,repeat 1):(2,repeat 2):(3,repeat 3):⊥ groupOn
```

So there is some kind of well-defined lazy semantics for this function. The puzzle I was interested in was defining an efficient implementation for these semantics.

# The Slow Case

The first approximation to a solution I could think of is the following:

```
groupOn :: Ord k => (a -> k) -> [a] -> [(k, [a])]
= Map.toList . Map.fromListWith (++) . map (\x -> (k x, [x])) groupOn k
```

In fact, if you don’t care about laziness, this is probably the best solution: it’s $\mathcal{O}(n \log n)$, it performs well (practically as well as asymptotically), and it has the expected results.

However, there are problems. Primarily this solution cares about ordering, which we don’t want. We want to emit the results in the same order that they were in the original list, and we don’t necessarily want to require an ordering on the elements (for the efficient solution we will relax this last constraint).

Instead, let’s implement our own “map” type that is inefficient, but more general.

```
type Map a b = [(a,b)]
insertWith :: Eq a => (b -> b -> b) -> a -> b -> Map a b -> Map a b
= [(k,v)]
insertWith f k v [] :xs)
insertWith f k v ((k',v')| k == k' = (k',f v v') : xs
| otherwise = (k',v') : insertWith f k v xs
groupOn :: Eq k => (a -> k) -> [a] -> [(k, [a])]
= foldr (uncurry (insertWith (++))) [] . map (\x -> (k x, [x])) groupOn k
```

The problem here is that it’s not lazy enough. `insertWith`

is strict in its last argument, which means that using `foldr`

doesn’t gain us anything laziness-wise.

There is some extra information we can use to drive the result: we know that the result will have keys that are in the same order as they appear in the list, with duplicates removed:

```
groupOn :: Eq k => (a -> k) -> [a] -> [(k, [a])]
= map _ ks
groupOn k xs where
= map k xs ks
```

From here, we can get what the values should be from each key by filtering the original list:

```
groupOn :: Eq k => (a -> k) -> [a] -> [(k,[a])]
= map (\k -> (k, filter ((k==) . key) xs)) (nub (map key xs)) groupOn key xs
```

Using a kind of Schwartzian transform yields the following slight improvement:

```
groupOn :: Eq k => (a -> k) -> [a] -> [(k,[a])]
= map (\k -> (k , map snd (filter ((k==) . fst) ks))) (nub (map fst ks))
groupOn key xs where
= map (\x -> (key x, x)) xs ks
```

But this traverses the same list multiple times unnecessarily. The problem is that we’re repeating a lot of work between `nub`

and the rest of the algorithm.

The following is much better:

```
groupOn :: Eq k => (a -> k) -> [a] -> [(k,[a])]
= go . map (\x -> (key x, x))
groupOn key where
= []
go [] :xs) = (k,x:map snd y) : go ys
go ((k,x)where
= partition ((k==).fst) xs (y,ys)
```

First, we perform the Schwartzian transform optimisation. The work of the algorithm is done in the `go`

helper. The idea is to filter out duplicates as we encounter them: when we encounter `(k,x)`

we can keep it immediately, but then we split the rest of the list into the components that have the same key as this element, and the ones that differ. The ones that have the same key can form the collection for this key, and those that differ are what we recurse on.

This partitioning also avoids re-traversing elements we know to be already accounted for in a previous group. I think that this is the most efficient (modulo some inlining and strictness improvements) algorithm that can do `groupOn`

with just an `Eq`

constraint.

# A Faster Version

The reason that the `groupOn`

above is slow is that every element returned has to traverse the entire rest of the list to remove duplicates. This is a classic pattern of quadratic behaviour: we can improve it by using the same trick as quick sort, by partitioning the list into lesser and greater elements on every call.

```
groupOnOrd :: Ord k => (a -> k) -> [a] -> [(k,[a])]
= go . map (\x -> (key x, x))
groupOnOrd key where
= []
go [] :xs) = (k,x:e) : go lt ++ go gt
go ((k,x)where
= foldr split ([],[],[]) xs
(e,lt,gt) @(k',y) ~(e,lt,gt) = case compare k' k of
split kyLT -> (e, ky:lt, gt)
EQ -> (y:e, lt, gt)
GT -> (e, lt, ky:gt)
```

While this is $\mathcal{O}(n \log n)$, and it does group elements, it also reorders the underlying list. Let’s fix that by tagging the incoming elements with their positions, and then using those positions to order them back into their original configuration:

```
groupOnOrd :: Ord k => (a -> k) -> [a] -> [(k,[a])]
= map (\(_,k,xs) -> (k,xs)) . go . zipWith (\i x -> (i, k x, x)) [0..]
groupOnOrd k where
= []
go [] :xs) = (i, k, x : e) : merge (go l) (go g)
go ((i, k, x)where
= foldr split ([],[],[]) xs
(e, l, g)
@(_,k',y) ~(e, l, g) = case compare k' k of
split kyLT -> (e , ky : l, g)
EQ -> (y:e, l, g)
GT -> (e , l, ky : g)
= gt
merge [] gt = lt
merge lt [] @(i,_,_):lt) (g@(j,_,_):gt)
merge (l| i <= j = l : merge lt (g:gt)
| otherwise = g : merge (l:lt) gt
```

This is close, but still not right. This isn’t yet *lazy*. The `merge`

function is strict in both arguments.

However, we have all the information we need to unshuffle the lists without having to inspect them. In `split`

, we know which direction we put each element: we can store that info without using indices.

```
groupOnOrd :: Ord k => (a -> k) -> [a] -> [(k,[a])]
= catMaybes . go . map (\x -> (k x, x))
groupOnOrd k where
= []
go [] :xs) = Just (k, x : e) : merge m (go l) (go g)
go ((k,x)where
= foldr split ([],[],[],[]) xs
(e, m, l, g)
@(k',y) ~(e, m, l, g) = case compare k' k of
split kyLT -> ( e, LT : m, ky : l, g)
EQ -> (y:e, EQ : m, l, g)
GT -> ( e, GT : m, l, ky : g)
= []
merge [] lt gt EQ : xs) lt gt = Nothing : merge xs lt gt
merge (LT : xs) (l:lt) gt = l : merge xs lt gt
merge (GT : xs) lt (g:gt) = g : merge xs lt gt merge (
```

What we generate here is a `[Ordering]`

: this list tells us the result of all the compare operations on the input list. Then, in `merge`

, we invert the action of `split`

, rebuilding the original list without inspecting either `lt`

or `gt`

.

And this solution works! It’s $\mathcal{O}(n \log n)$, and fully lazy.

```
>>> map fst . groupOnOrd id $ [1..]
1..]
[
>>> groupOnOrd id $ cycle [1,2,3]
1,repeat 1):(2,repeat 2):(3,repeat 3):⊥
(
>>> groupOnOrd (`rem` 3) [1..]
1,[1,4..]):(2,[2,5..]):(0,[3,6..]):⊥ (
```

The finished version of these two functions, along with some benchmarks, is available here.