A *union-find* data
structure
(also known as a *disjoint set* data structure) keeps track of a
*collection of disjoint sets*, typically with elements drawn from
\(\{0, \dots, n-1\}\). For example, we might have the sets

\(\{1,3\}, \{0, 4, 2\}, \{5, 6, 7\}\)

A union-find structure must support three basic operations:

We can \(\mathit{create}\) a union-find structure with \(n\) singleton sets \(\{0\}\) through \(\{n-1\}\). (Alternatively, we could support two operations: creating an empty union-find structure, and adding a new singleton set; occasionally this more fine-grained approach is useful, but we will stick with the simpler \(\mathit{create}\) API for now.)

We can \(\mathit{find}\) a given \(x \in \{0, \dots, n-1\}\), returning some sort of “name” for the set \(x\) is in. It doesn’t matter what these names are; the only thing that matters is that for any \(x\) and \(y\), \(\mathit{find}(x) = \mathit{find}(y)\) if and only if \(x\) and \(y\) are in the same set. The most important application of \(\mathit{find}\) is therefore to check whether two given elements are in the same set or not.

We can \(\mathit{union}\) two elements, so the sets that contain them become one set. For example, if we \(\mathit{union}(2,5)\) then we would have

\(\{1,3\}, \{0, 4, 2, 5, 6, 7\}\)

Note that \(\mathit{union}\) is a one-way operation: once two sets have been unioned together, there’s no way to split them apart again. (If both merging and splitting are required, one can use a link/cut tree, which is very cool—and possibly something I will write about in the future—but much more complex.) However, these three operations are enough for union-find structures to have a large number of interesting applications!

In addition, we can annotate each set with a value taken from some commutative semigroup. When creating a new union-find structure, we must specify the starting value for each singleton set; when unioning two sets, we combine their annotations via the semigroup operation.

- For example, we could annotate each set with its
*size*; singleton sets always start out with size 1, and every time we union two sets we add their sizes. - We could also annotate each set with the sum, product, maximum, or minumum of all its elements.
- Of course there are many more exotic examples as well.

We typically use a commutative semigroup, as in the examples above;
this guarantees that a given set always has a single well-defined
annotation value, regardless of the sequence of union-find operations
that were used to create it. However, we can actually use any binary
operation at all (*i.e.* any *magma*), in which case the annotations
on a set may reflect the precise tree of calls to \(\mathit{union}\) that were
used to construct it; this can occasionally be useful.

For example, we could annotate each set with a list of values, and combine annotations using list concatenation; the order of elements in the list associated to a given set will depend on the order of arguments to \(\mathit{union}\).

We could also annotate each set with a binary tree storing values at the leaves. Each singleton set is annotated with a single leaf; to combine two trees we create a new branch node with the two trees as its children. Then each set ends up annotated with the precise tree of calls to \(\mathit{union}\) that were used to create it.

My implementation is based on one by Kwang Yul
Seo,
but I have modified it quite a bit. The code is also available in my
`comprog-hs`

repository. This
blog post is not intended to be a comprehensive union-find tutorial,
but I will explain some things as we go.

```
{-# LANGUAGE RecordWildCards #-}
module UnionFind where
import Control.Monad (when)
import Control.Monad.ST
import Data.Array.ST
```

Let’s start with the definition of the `UnionFind`

type itself.
`UnionFind`

has two type parameters: `s`

is a phantom type parameter
used to limit the scope to a given `ST`

computation; `m`

is the type
of the arbitrary annotations. Note that the elements are also
sometimes called “nodes”, since, as we will see, they are organized
into trees.

```
type Node = Int
data UnionFind s m = UnionFind {
```

The basic idea is to maintain three mappings:

- First, each element is mapped to a
*parent*(another element). There are no cycles, except that some elements can be their own parent. This means that the elements form a*forest*of rooted trees, with the self-parenting elements as roots. We store the parent mapping as an`STUArray`

(see here for another post where we used`STUArray`

) for efficiency.

` parent :: !(STUArray s Node Node),`

Each element is also mapped to a

*size*. We maintain the invariant that for any element which is a root (*i.e.*any element which is its own parent), we store the size of the tree rooted at that element. The size associated to other, non-root elements does not matter.(Many implementations store the

*height*of each tree instead of the size, but it does not make much practical difference, and the size seems more generally useful.)

` sz :: !(STUArray s Node Int),`

- Finally, we map each element to a custom annotation value; again, we only care about the annotation values for root nodes.

` ann :: !(STArray s Node m) }`

To \(\mathit{create}\) a new union-find structure, we need a size and a
function mapping each element to an initial annotation value. Every
element starts as its own parent, with a size of 1. For convenience,
we can also make a variant of `createWith`

that gives every element
the same constant annotation value.

```
createWith :: Int -> (Node -> m) -> ST s (UnionFind s m)
=
createWith n m UnionFind
<$> newListArray (0, n - 1) [0 .. n - 1] -- Every node is its own parent
<*> newArray (0, n - 1) 1 -- Every node has size 1
<*> newListArray (0, n - 1) (map m [0 .. n - 1])
create :: Int -> m -> ST s (UnionFind s m)
= createWith n (const m) create n m
```

To perform a \(\mathit{find}\) operation, we keep following *parent*
references up the tree until reaching a root. We can also do a cool
optimization known as *path compression*: after finding a
root, we can directly update the parent of every node along the path
we just traversed to be the root. This means \(\mathit{find}\) can be very
efficient, since it tends to create trees that are extremely wide and
shallow.

```
find :: UnionFind s m -> Node -> ST s Node
@(UnionFind {..}) x = do
find uf<- readArray parent x
p if p /= x
then do
<- find uf p
r
writeArray parent x rpure r
else pure x
connected :: UnionFind s m -> Node -> Node -> ST s Bool
= (==) <$> find uf x <*> find uf y connected uf x y
```

Finally, to implement \(\mathit{union}\), we find the roots of the given nodes; if they are not the same we make the root with the smaller tree the child of the other root, combining sizes and annotations as appropriate.

```
union :: Semigroup m => UnionFind s m -> Node -> Node -> ST s ()
@(UnionFind {..}) x y = do
union uf<- find uf x
x <- find uf y
y /= y) $ do
when (x <- readArray sz x
sx <- readArray sz y
sy <- readArray ann x
mx <- readArray ann y
my if sx < sy
then do
writeArray parent x y+ sy)
writeArray sz y (sx <> my)
writeArray ann y (mx else do
writeArray parent y x+ sy)
writeArray sz x (sx <> my) writeArray ann x (mx
```

Note the trick of writing `x <- find uf x`

: this looks kind of like an
imperative statement that updates the value of a mutable variable `x`

,
but really it just makes a new variable `x`

which shadows the old
one.

Finally, a few utility functions. First, one to get the size of the set containing a given node:

```
size :: UnionFind s m -> Node -> ST s Int
@(UnionFind {..}) x = do
size uf<- find uf x
x readArray sz x
```

Also, we can provide functions to update and fetch the custom annotation value associated to the set containing a given node.

```
updateAnn :: Semigroup m => UnionFind s m -> Node -> m -> ST s ()
@(UnionFind {..}) x m = do
updateAnn uf<- find uf x
x <- readArray ann x
old <> m)
writeArray ann x (old -- We could use modifyArray above, but the version of the standard library
-- installed on Kattis doesn't have it
getAnn :: UnionFind s m -> Node -> ST s m
@(UnionFind {..}) x = do
getAnn uf<- find uf x
x readArray ann x
```

Here are a couple of problems I challenge you to solve for next time:

**tl;dr**: a fix to the `MonadRandom`

package may cause `fromListMay`

and related functions to *extremely rarely* output different results than
they used to. This could only possibly affect anyone who is using
fixed seed(s) to generate random values and is depending on the
specific values being produced, *e.g.* a unit test where you use a
specific seed and test that you get a specific result. Do you think
this should be a major or minor version bump?

Since 2013 I have been the maintainer of
`MonadRandom`

,
which defines a monad and monad transformer for generating random
values, along with a number of related utilities.

Recently, Toni Dietze pointed out a rare
situation that could cause the `fromListMay`

function to
crash (as well as
the other functions which depend on it: `fromList`

, `weighted`

,
`weightedMay`

, `uniform`

, and `uniformMay`

). This function is
supposed to draw a weighted random sample from a list of values
decorated with weights. I’m not going to explain the details of the
issue here; suffice it to say that it has to do with conversions
between `Rational`

(the type of the weights) and `Double`

(the type
that was being used internally for generating random numbers).

Even though this could only happen in rare and/or strange
circumstances, fixing it definitely seemed like the right thing to
do. After a bit of discussion, Toni came up with a good suggestion
for a fix: we should no longer use `Double`

internally for generating
random numbers, but rather `Word64`

, which avoids conversion and
rounding issues.

In fact, `Word64`

is already used internally in the generation of
random `Double`

values, so we can emulate the behavior of the `Double`

instance (which was slightly
tricky
to figure out)
so that we make exactly the same random choices as before, but without
actually converting to `Double`

.

…well, not *exactly* the same random choices as before, and therein
lies the rub! If `fromListMay`

happens to pick a random value which
is extremely close to a boundary between choices, it’s possible that
the value will fall on one side of the boundary when using exact
calculations with `Word64`

and `Rational`

, whereas before it would
have fallen on the other side of the boundary after converting to
`Double`

due to rounding. In other words, it will output the
same results *almost all the time*, but for a list of \(n\) weighted
choices there is something like an \(n/2^{64}\) chance (or less) that
any given random choice will be different from what it used to be. I
have never observed this happening in my tests, and indeed, I do not
expect to ever observe it! If we generated one billion random samples
per second continuously for a thousand years, we might expect to see
it happen once or twice. I am not even sure how to engineer a test
scenario to force it to happen, because we would have to pick an
initial PRNG seed that forces a certain `Word64`

value to be
generated.

Technically, a function exported by `MonadRandom`

has changed
behavior, so according to the Haskell PVP
specification this should be a major
version bump (*i.e.* `0.6`

to `0.7`

).Actually, I am not even
100% clear on this. The decision
tree on the PVP page says
that changing the *behavior* of an exported function necessitates a
major version bump; but the actual
specification does not
refer to *behavior* at all—as I read it, it is exclusively concerned
with API compatibility, *i.e.* whether things will still compile.

But
there seem to be some good arguments for doing just a minor version
bump (*i.e.* `0.6`

to `0.6.1`

).

Arguments in favor of a minor version bump:

A major version bump would cause a lot of (probably unnecessary) breakage!

`MonadRandom`

has 149 direct reverse dependencies, and about 3500 distinct transitive reverse dependencies. Forcing all those packages to update their upper bound on`MonadRandom`

would be a lot of churn.What exactly constitutes the “behavior” of a function to generate random values? It depends on your point of view. If we view the function as a pure mathematical function which takes a PRNG state as input and produces some value as output, then its behavior is defined precisely by which outputs it returns for which input seeds, and its behavior has changed. However, if we think of it in more effectful terms, we could say its “behavior” is just to output random values according to a certain distribution, in which case its behavior has

*not*changed.It’s extremely unlikely that this change will cause any breakage; moreover, as argued by Boyd Stephen Smith, anyone who cares enough about reproducibility to be relying on specific outputs for specific seeds is probably already pinning all their package versions.

Arguments in favor of a major version bump:

It’s what the PVP specifies; what’s the point of having a specification if we don’t follow it?

In the unlikely event that this change

*does*cause any breakage, it could be extremely difficult for package maintainers to track down. If the behavior of a random generation function completely changes, the source of the issue is obvious. But if it only changes for very rare inputs, you might reasonably think the problem is something else. A major version bump will force maintainers to read the changelog for`MonadRandom`

and assess whether this is a change that could possibly affect them.

So, do you have opinions on this? Would the release affect you one way or the other? Feel free to leave a comment here, or send me an email with your thoughts. Note there has already been a bit of discussion on Mastodon as well.

Way back in 2012 I took over maintainership of the `BlogLiterately`

tool from Robert
Greayer, its initial author. I used it for many years to post to my
Wordpress blog, added a
bunch
of
features,
solved some fun
bugs,
and created the accompanying `BlogLiterately-diagrams`

plugin
for embedding diagrams code in blog
posts. However, now that I have fled Wordpress and rebuilt my blog
with hakyll, I don’t use
`BlogLiterately`

any more (there is even a `diagrams-pandoc`

package
which does the same thing `BlogLiterately-diagrams`

used to do). So,
as of today I am officially declaring `BlogLiterately`

unsupported.

The fact is, I haven’t actually updated `BlogLiterately`

since March
of last year. It currently only builds on GHC 9.4 or older, and no one
has complained, which I take as strong evidence that no one else is
using it either! However, if anyone out there is actually using it,
and would like to take over as maintainer, I would be very happy to
pass it along to you.

I do plan to continue maintaining
`HaXml`

and
`haxr`

, at least for now;
unlike `BlogLiterately`

, I know they are still in use, especially
`HaXml`

. However, `BlogLiterately`

was really the only reason I cared
about these packages personally, so I would be happy to pass them
along as well; please get in touch if you would be willing to take
over maintaining one or both packages.

Recently, as part of a larger project, I wanted to define decidable equality for an indexed data type in Agda. I struggled quite a bit to figure out the right way to encode it to make Agda happy, and wasn’t able to find much help online, so I’m recording the results here.

The **tl;dr** is to use mutual recursion to define the indexed data
type along with a sigma type that hides the index, and to use the
sigma type in any recursive positions where we don’t care about the
index! Read on for more motivation and details (and wrong turns I
took along the way).

This post is literate Agda; you can download it here if you want to play along. I tested everything here with Agda version 2.6.4.3 and version 2.0 of the standard library.

First, some imports and a module declaration. Note that the entire
development is parameterized by some abstract set `B`

of base types,
which must have decidable equality.

```
open import Data.Product using (Σ ; _×_ ; _,_ ; -,_ ; proj₁ ; proj₂)
open import Data.Product.Properties using (≡-dec)
open import Function using (_∘_)
open import Relation.Binary using (DecidableEquality)
open import Relation.Binary.PropositionalEquality using (_≡_ ; refl)
open import Relation.Nullary.Decidable using (yes; no; Dec)
module OneLevelTypesIndexed (B : Set) (≟B : DecidableEquality B) where
```

We’ll work with a simple type system containing base types, function types, and some distinguished type constructor □. So far, this is just to give some context; it is not the final version of the code we will end up with, so we stick it in a local module so it won’t end up in the top-level namespace.

```
module Unindexed where
data Ty : Set where
: B → Ty
base _⇒_ : Ty → Ty → Ty
_ : Ty → Ty □
```

For example, if \(X\) and \(Y\) are base types, then we could write down a type like \(\square ((\square \square X \to Y) \to \square Y)\):

```
infixr 2 _⇒_
infix 30 □_
postulate
: B
BX BY
: Ty
X = base BX
X : Ty
Y = base BY
Y
: Ty
example = □ ((□ □ X ⇒ Y) ⇒ □ Y) example
```

However, for reasons that would take us too far afield in this blog
post, I *don’t* want to allow immediately nested boxes, like \(\square \square X\). We can still have multiple boxes in a type, and even
boxes nested inside of other boxes, as long as there is at least one
arrow in between. In other words, I only want to rule out boxes
immediately applied to another type with an outermost box. So we
don’t want to allow the example type given above (since it contains
\(\square \square X\)), but, for example, \(\square ((\square X \to Y) \to \square Y)\) would be OK.

How can we encode this invariant so it holds by construction? One way would be to have two mutually recursive data types, like so:

```
module Mutual where
data Ty : Set
data UTy : Set
data Ty where
_ : UTy → Ty
□_ : UTy → Ty
∙
data UTy where
: B → UTy
base _⇒_ : Ty → Ty → UTy
```

`UTy`

consists of types which have no top-level box; the constructors
of `Ty`

just inject `UTy`

into `Ty`

by adding either one or zero
boxes. This works, and defining decidable equality for `Ty`

and `UTy`

is relatively straightforward (again by mutual recursion). However,
it seemed to me that having to deal with `Ty`

and `UTy`

everywhere
through the rest of the development was probably going to be super
annoying.

The other option would be to index `Ty`

by values indicating whether a
type has zero or one top-level boxes; then we can use the indices to
enforce the appropriate rules. First, we define a data type `Boxity`

to act as the index for `Ty`

, and show that it has decidable equality:

```
data Boxity : Set where
: Boxity
[0] : Boxity
[1]
: DecidableEquality Boxity
Boxity-≟ = yes refl
Boxity-≟ [0] [0] = no λ ()
Boxity-≟ [0] [1] = no λ ()
Boxity-≟ [1] [0] = yes refl Boxity-≟ [1] [1]
```

My first attempt to write down a version of `Ty`

indexed by `Boxity`

looked like this:

```
module IndexedTry1 where
data Ty : Boxity → Set where
: B → Ty [0]
base _⇒_ : {b₁ b₂ : Boxity} → Ty b₁ → Ty b₂ → Ty [0]
_ : Ty [0] → Ty [1] □
```

`base`

always introduces a type with no top-level box; the `□`

constructor requires a type with no top-level box, and produces a type
with one (this is what ensures we cannot nest boxes); and the arrow
constructor does not care how many boxes its arguments have, but
constructs a type with no top-level box.

This is logically correct, but I found it very difficult to work with.
The sticking point for me was injectivity of the arrow constructor.
When defining decidable equality we need to prove lemmas that each of
the constructors are injective, but I was not even able to write down
the *type* of injectivity for `_⇒_`

. We would want something like this:

```
-inj :
⇒: Boxity}
{bσ₁ bσ₂ bτ₁ bτ₂ : Ty bσ₁} {σ₂ : Ty bσ₂} {τ₁ : Ty bτ₁} {τ₂ : Ty bτ₂} →
{σ₁
(σ₁ ⇒ σ₂) ≡ (τ₁ ⇒ τ₂) → (σ₁ ≡ τ₁) × (σ₂ ≡ τ₂)
```

but this does not even typecheck! The problem is that, for example,
`σ₁`

and `τ₁`

have different types, so the equality proposition `σ₁ ≡ τ₁`

is not well-typed.

At this point I tried turning to heterogeneous equality, but it didn’t seem to help. I won’t record here all the things I tried, but the same issues seemed to persist, just pushed around to different places (for example, I was not able to pattern-match on witnesses of heterogeneous equality because of types that didn’t match).

At ICFP last week I asked Jesper Cockx
for advice,which felt a bit like asking Rory McIlroy to give some
tips on your mini-golf game

and he suggested trying to prove
decidable equality for the sigma type pairing an index with a type
having that index, like this:

```
: Set
ΣTy = Σ Boxity Ty ΣTy
```

This turned out to be the key idea, but it still took me a long time
to figure out the right way to make it work. Given the above
definitions, if we go ahead and try to define decidable equality for
`ΣTy`

, injectivity of the arrow constructor is still a problem.

After days of banging my head against this off and on, I finally
realized that the way to solve this is to define `Ty`

and `ΣTy`

by
mutual recursion: the arrow constructor should just take two `ΣTy`

arguments! This perfectly captures the idea that we *don’t care*
about the indices of the arrow constructor’s argument types, so we
hide them by bundling them up in a sigma type.

```
: Set
ΣTy data Ty : Boxity → Set
= Σ Boxity Ty
ΣTy
data Ty where
_ : Ty [0] → Ty [1]
□: B → Ty [0]
base _⇒_ : ΣTy → ΣTy → Ty [0]
infixr 2 _⇒_
infix 30 □_
```

Now we’re cooking! We now make quick work of the required injectivity
lemmas, which all go through trivially by matching on `refl`

:

```
: {τ₁ τ₂ : Ty [0]} → (□ τ₁ ≡ □ τ₂) → (τ₁ ≡ τ₂)
□-inj = refl
□-inj refl
: {b₁ b₂ : B} → base b₁ ≡ base b₂ → b₁ ≡ b₂
base-inj = refl
base-inj refl
: {σ₁ σ₂ τ₁ τ₂ : ΣTy} → (σ₁ ⇒ σ₂) ≡ (τ₁ ⇒ τ₂) → (σ₁ ≡ τ₁) × (σ₂ ≡ τ₂)
⇒-inj = refl , refl ⇒-inj refl
```

Notice how the type of `⇒-inj`

is now perfectly fine: we just have a
bunch of `ΣTy`

values that hide their indices, so we can talk about
propositional equality between them with no trouble.

Finally, we can define decidable equality for `Ty`

and `ΣTy`

by mutual
recursion.

```
: DecidableEquality ΣTy
ΣTy-≟
{-# TERMINATING #-}
: ∀ {b} → DecidableEquality (Ty b) Ty-≟
```

Sadly, I had to reassure Agda that the definition of `Ty-≟`

is terminating—more on this later.

To define `ΣTy-≟`

we can just use a lemma from
`Data.Product.Properties`

which derives decidable equality for a sigma
type from decidable equality for both components.

`= ≡-dec Boxity-≟ Ty-≟ ΣTy-≟ `

The only thing left is to define decidable equality for any two values
of type `Ty b`

(given a specific boxity `b`

), making use of our
injectivity lemmas; now that we have the right definitions, this falls
out straightforwardly.

```
(□ σ) (□ τ) with Ty-≟ σ τ
Ty-≟ ... | no σ≢τ = no (σ≢τ ∘ □-inj)
... | yes refl = yes refl
(base x) (base y) with ≟B x y
Ty-≟ ... | no x≢y = no (x≢y ∘ base-inj)
... | yes refl = yes refl
(σ₁ ⇒ σ₂) (τ₁ ⇒ τ₂) with ΣTy-≟ σ₁ τ₁ | ΣTy-≟ σ₂ τ₂
Ty-≟ ... | no σ₁≢τ₁ | _ = no (σ₁≢τ₁ ∘ proj₁ ∘ ⇒-inj)
... | yes _ | no σ₂≢τ₂ = no (σ₂≢τ₂ ∘ proj₂ ∘ ⇒-inj)
... | yes refl | yes refl = yes refl
(base _) (_ ⇒ _) = no λ ()
Ty-≟ (_ ⇒ _) (base _) = no λ () Ty-≟
```

First, the one remaining infelicity is that Agda could not tell that
`Ty-≟`

is terminating. I am not entirely sure why, but I think it may
be that the way the recursion works is just too convoluted for it to
analyze properly: `Ty-≟`

calls `ΣTy-≟`

on structural subterms of its
inputs, but then `ΣTy-≟`

works by providing `Ty-≟`

*as a higher-order
parameter* to `≡-dec`

. If you look at the definition of `≡-dec`

, all
it does is call its function parameters on structural subterms of its
input, so everything should be nicely terminating, but I guess I am
not surprised that Agda is not able to figure this out. If anyone has
suggestions on how to make this pass the termination checker without
using a `TERMINATING`

pragma, I would love to hear it!

As a final aside, I note that converting back and forth between `Ty`

(with `ΣTy`

arguments to the arrow constructor) and `IndexedTry1.Ty`

(with expanded-out `Boxity`

and `Ty`

arguments to arrow) is trivial:

```
: {b : Boxity} → Ty b → IndexedTry1.Ty b
Ty→Ty1 (□ σ) = IndexedTry1.□ (Ty→Ty1 σ)
Ty→Ty1 (base x) = IndexedTry1.base x
Ty→Ty1 ((b₁ , σ₁) ⇒ (b₂ , σ₂)) = (Ty→Ty1 σ₁) IndexedTry1.⇒ (Ty→Ty1 σ₂)
Ty→Ty1
: {b : Boxity} → IndexedTry1.Ty b → Ty b
Ty1→Ty (IndexedTry1.base x) = base x
Ty1→Ty (σ₁ IndexedTry1.⇒ σ₂) = -, (Ty1→Ty σ₁) ⇒ -, (Ty1→Ty σ₂)
Ty1→Ty (IndexedTry1.□ σ) = □ (Ty1→Ty σ) Ty1→Ty
```

I expect it is also trivial to prove this is an isomorphism, though I’m not particularly motivated to do it. The point is that, as anyone who has spent any time proving things with proof assistants knows, two types can be completely isomorphic, and yet one can be vastly easier to work with than the other in certain contexts. Often when I’m trying to prove something in Agda it feels like at least half the battle is just coming up with the right representation that makes the proofs go through easily.

In a previous
post
I discussed the first half of my solution to Factor-Full
Tree. In this post,
I will demonstrate how to *decompose a tree into disjoint paths*.
Technically, we should clarify that we are looking for *directed*
paths in a rooted tree, that is, paths that only proceed down the
tree. One could also ask about decomposing an unrooted tree into
disjoint undirected paths; I haven’t thought about how to do that in
general but intuitively I expect it is not too much more difficult.

For
this particular problem, we want to decompose a tree into
*maximum-length* paths (*i.e.* we start by taking the longest possible
path, then take the longest path from what remains, and so on); I will call
this the *max-chain decomposition* (I don’t know if there is a
standard term). However, there are other types of path
decomposition, such as heavy-light decomposition, so we will try to
keep the decomposition code somewhat generic.

This post is literate Haskell; you can find the source code on GitHub. We begin with some language pragmas and imports.

```
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TupleSections #-}
module TreeDecomposition where
import Control.Arrow ((>>>), (***))
import Data.Bifunctor (second)
import Data.ByteString.Lazy.Char8 (ByteString)
import Data.ByteString.Lazy.Char8 qualified as BS
import Data.List (sortBy)
import Data.List.NonEmpty (NonEmpty)
import Data.List.NonEmpty qualified as NE
import Data.Map (Map, (!), (!?))
import Data.Map qualified as M
import Data.Ord (Down(..), comparing)
import Data.Tree (Tree(..), foldTree)
import Data.Tuple (swap)
import ScannerBS
```

Remember, our goal is to split up a tree into a collection of linear paths; that is, in general, something like this:

What do we need in order to specify a decomposition of a
tree into disjoint paths this way? Really, all we need is to choose *at most
one linked child* for each node. In other words, at every node we can
choose to continue the current path into a single child node (in which
case all the other children will start their own new paths), or we
could choose to terminate the current path (in which case every child
will be the start of its own new path). We can represent such a
choice with a function of type

`type SubtreeSelector a = a -> [Tree a] -> Maybe (Tree a, [Tree a])`

which takes as input the value at a node and the list of all the
subtrees, and possibly returns a selected subtree along with the list of remaining
subtrees.Of course, there is nothing in the
type that actually requires a `SubtreeSelector`

to return one of the
trees from its input paired with the rest, but nothing we will do
depends on this being true. In fact, I expect there may be some
interesting algorithms obtainable by running a “path decomposition”
with a “selector” function that actually makes up new trees instead of just
selecting one, similar to the `chop`

function.

Given such a subtree selection function, a generic path decomposition
function will then take a tree and turn it into a list of non-empty
paths:We could also imagine wanting information about the parent of each
path, and a mapping from tree nodes to some kind of path ID, but we
will keep things simple for now.

`pathDecomposition :: SubtreeSelector a -> Tree a -> [NonEmpty a]`

Implementing `pathDecomposition`

is a nice exercise; you might like to
try it yourself! You can find my implementation at the end of this
blog post.

Now, let’s use our generic path decomposition to implement a max-chain
decomposition. At each node we want to select the *tallest* subtree;
in order to do this efficiently, we can first annotate each tree node with
its height, via a straightforward tree fold:

```
type Height = Int
labelHeight :: Tree a -> Tree (Height, a)
= foldTree node
labelHeight where
= case ts of
node a ts -> Node (0, a) []
[] -> Node (1 + maximum (map (fst . rootLabel) ts), a) ts _
```

Our subtree selection function can now select the subtree with the
largest `Height`

annotation. Instead of implementing this directly,
we might as well make a generic function for selecting the “best”
element from a list (we will reuse it later):

```
selectMaxBy :: (a -> a -> Ordering) -> [a] -> Maybe (a, [a])
= Nothing
selectMaxBy _ [] : as) = case selectMaxBy cmp as of
selectMaxBy cmp (a Nothing -> Just (a, [])
Just (b, bs) -> case cmp a b of
LT -> Just (b, a : bs)
-> Just (a, b : bs) _
```

We can now put the pieces together to implement max-chain
decomposition. We first label the tree by height, then do a path
decomposition that selects the tallest subtree at each node. We leave
the height annotations in the final output since they might be
useful—for example, we can tell how long each path is just by
looking at the `Height`

annotation on the first element. If we don’t
need them we can easily get rid of them later. We also sort by
descending `Height`

, since getting the longest chains first was kind
of the whole point.

```
maxChainDecomposition :: Tree a -> [NonEmpty (Height, a)]
=
maxChainDecomposition >>>
labelHeight const (selectMaxBy (comparing (fst . rootLabel)))) >>>
pathDecomposition (Down . fst . NE.head)) sortBy (comparing (
```

To flesh this out into a full solution to Factor-Full
Tree, after
computing the chain decomposition we need to assign prime factors to
the chains. From those, we can compute the value for each node if we
know which chain it is in and the value of its parent. To this end,
we will need one more function which computes a `Map`

recording the
parent of each node in a tree. Note that if we already know all the
edges in a given edge list are oriented the same way, we can build
this much more simply as *e.g.* `map swap >>> M.fromList`

; but when
(as in general) we don’t know which way the edges should be oriented
first, we might as well first build a `Tree a`

via DFS with
`edgesToTree`

and then construct the `parentMap`

like this afterwards.

```
parentMap :: Ord a => Tree a -> Map a a
= foldTree node >>> snd
parentMap where
node :: Ord a => a -> [(a, Map a a)] -> (a, Map a a)
= (a, M.fromList (map (,a) as) <> mconcat ms)
node a b where
= unzip b (as, ms)
```

Finally, we can solve Factor-Full tree. Note that some code from my
previous blog
post
is needed as well, and is included at the end of the post for
completeness. Once we compute the max chain decomposition and the
prime factor for each node, we use a lazy recursive
`Map`

to compute the value assigned to each node.

```
solve :: TC -> [Int]
TC{..} = M.elems assignment
solve where
-- Build the tree and compute its parent map
= edgesToTree Node edges 1
t = parentMap t
parent
-- Compute the max chain decomposition, and use it to assign a prime factor
-- to each non-root node
paths :: [[Node]]
= map (NE.toList . fmap snd) $ maxChainDecomposition t
paths
factor :: Map Node Int
= M.fromList . concat $ zipWith (\p -> map (,p)) primes paths
factor
-- Compute an assignment of each node to a value, using a lazy map
assignment :: Map Node Int
= M.fromList $ (1,1) : [(v, factor!v * assignment!(parent!v)) | v <- [2..n]] assignment
```

For an explanation of this code for `primes`

, see this old blog post.

```
primes :: [Int]
= 2 : sieve primes [3 ..]
primes where
: ps) xs =
sieve (p let (h, t) = span (< p * p) xs
in h ++ sieve ps (filter ((/= 0) . (`mod` p)) t)
```

We can easily use our generic path decomposition to compute a heavy-light decomposition as well:

```
type Size = Int
labelSize :: Tree a -> Tree (Size, a)
= foldTree $ \a ts -> Node (1 + sum (map (fst . rootLabel) ts), a) ts
labelSize
heavyLightDecomposition :: Tree a -> [NonEmpty (Size, a)]
=
heavyLightDecomposition >>>
labelSize const (selectMaxBy (comparing (fst . rootLabel)))) pathDecomposition (
```

I plan to write about this in a future post.

Here’s my implementation of `pathDecomposition`

; how did you do?

```
= go
pathDecomposition select where
= selectPath select >>> second (concatMap go) >>> uncurry (:)
go
selectPath :: SubtreeSelector a -> Tree a -> (NonEmpty a, [Tree a])
= go
selectPath select where
Node a ts) = case select a ts of
go (Nothing -> (NE.singleton a, ts)
Just (t, ts') -> ((a NE.<|) *** (ts' ++)) (go t)
```

We also include some input parsing and tree-building code from last time.

```
main :: IO ()
= BS.interact $ runScanner tc >>> solve >>> map (show >>> BS.pack) >>> BS.unwords
main
type Node = Int
data TC = TC { n :: Int, edges :: [(Node, Node)] }
deriving (Eq, Show)
tc :: Scanner TC
= do
tc <- int
n <- (n - 1) >< pair int int
edges return TC{..}
edgesToMap :: Ord a => [(a, a)] -> Map a [a]
= concatMap (\p -> [p, swap p]) >>> dirEdgesToMap
edgesToMap
dirEdgesToMap :: Ord a => [(a, a)] -> Map a [a]
= map (second (: [])) >>> M.fromListWith (++)
dirEdgesToMap
mapToTree :: Ord a => (a -> [b] -> b) -> Map a [a] -> a -> b
= dfs root root
mapToTree nd m root where
= nd root (maybe [] (map (dfs root) . filter (/= parent)) (m !? root))
dfs parent root
edgesToTree :: Ord a => (a -> [b] -> b) -> [(a, a)] -> a -> b
= mapToTree nd . edgesToMap edgesToTree nd
```

**tl;dr**: if you appreciate my past or ongoing contributions to the
Haskell community, please consider helping me get to ICFP this year by donating
via my ko-fi page!

Working at a small liberal arts institution has some tremendous benefits (close interaction with motivated students, freedom to pursue the projects I want rather than jump through a bunch of hoops to get tenure, fantastic colleagues), and I love my job. But there are also downsides; the biggest ones for me are the difficulty of securing enough travel funding, and, relatedly, the difficulty of cultivating and maintaining collaborations.

Last year I was very grateful for people’s generosity in helping me get to Seattle. I am planning to again attend ICFP in Milan this September; this time I will even bring some students along. I have once again secured some funding from my institution, but it will not be enough to cover all the expenses.

So, if you have been positively impacted by my
contributions to the Haskell community (*e.g.* CIS
194, the
Typeclassopedia,
diagrams,
split,
MonadRandom, burrito
metaphors…)
and/or would like to support my ongoing work (competitive programming
in
Haskell,
swarm,
disco, ongoing package
maintenance…), and are
able to express that appreciation or support with a donation of any
size to help me get to ICFP, I would really appreciate it!

Thank you, friends — I hope to see many people in Milan! Next up: I will soon publish another post about tree path decomposition!

Lately I’ve been thinking about representing *eventually constant
streams* in Haskell. An eventually constant stream is an infinite
stream which eventually, after some finite prefix, starts repeating
the same value forever. For example,

\(6, 8, 2, 9, 3, 1, 1, 1, 1, \dots\)

There are many things we can do in a decidable way with eventually constant streams that we can’t do with infinite streams in general—for example, test them for equality.

This is a work in progress. I only have one specific use case in mind (infinite-precision two’s complement arithmetic, explained at the end of the post), so I would love to hear of other potential use cases, or any other feedback. Depending on the feedback I may eventually turn this into a package on Hackage.

This blog post is typeset from a literate Haskell file; if you want to play along you can download the source from GitHub.

`River`

typeSome preliminaries:

```
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ViewPatterns #-}
module River where
import Data.Monoid (All (..), Any (..))
import Data.Semigroup (Max (..), Min (..))
import Prelude hiding (all, and, any, drop, foldMap, maximum, minimum, or, repeat, take, zipWith, (!!))
import Prelude qualified as P
```

Now let’s get to the main definition. A value of type `River a`

is
either a constant `C a`

, representing an infinite stream of copies of
`a`

, or a `Cons`

with an `a`

and a `River a`

.

```
data River a = C !a | Cons !a !(River a)
deriving Functor
```

I call this a `River`

since “all `River`

s flow to the `C`

”!

The strictness annotations on the `a`

values just seem like a good
idea in general. The strictness annotation on the `River a`

tail,
however, is more interesting: it’s there to rule out infinite streamsAlthough the strictness annotation on the `River a`

is semantically correct, I could imagine not wanting it there for performance reasons; I’d be happy to hear any feedback on this point.

constructed using only `Cons`

, such as `flipflop = Cons 0 (Cons 1 flipflop)`

. In
other words, the only way to make a non-bottom value of type `Stream a`

is
to have a finite sequence of `Cons`

finally terminated by `C`

.

We need to be a bit careful here, since there are multiple ways to
represent streams which are semantically supposed to be the same. For
example, `Cons 1 (Cons 1 (C 1))`

and `C 1`

both represent an infinite stream of
all `1`

’s. In general, we have the law

`C a === Cons a (C a)`

,

and want to make sure that any functions we write respect this
It would be interesting to try implementing rivers as a higher inductive type, say, in Cubical Agda.

equivalence, *i.e.* do not distinguish between such values. This is
the reason I did not derive an `Eq`

instance; we will have to write
our own.

We can partially solve this problem with a *bidirectional pattern
synonym*:

```
expand :: River a -> River a
C a) = Cons a (C a)
expand (= as
expand as
infixr 5 :::
pattern (:::) :: a -> River a -> River a
pattern (:::) a as <- (expand -> Cons a as)
where
::: as = Cons a as
a
{-# COMPLETE (:::) #-}
```

Matching with the pattern `(a ::: as)`

uses a *view pattern*
to potentially expand a `C`

one step into a `Cons`

, so that we can
pretend all `River`

values are always constructed with `(:::)`

.
In the other direction, `(:::)`

merely constructs a `Cons`

.

We mark `(:::)`

as `COMPLETE`

on its own since it is, in fact,
sufficient to handle every possible input of type `River`

. However,
in order to obtain terminating algorithms we will often include one or
more special cases for `C`

.

As an alternative, we could use a variant pattern synonym:

```
infixr 5 ::=
pattern (::=) :: Eq a => a -> River a -> River a
pattern (::=) a as <- (expand -> Cons a as)
where
::= C a | a' == a = C a
a' ::= as = Cons a as
a
{-# COMPLETE (::=) #-}
```

As compared to `(:::)`

, this has an extra `Eq a`

constraint: when we
construct a `River`

with `(::=)`

, it checks to see whether we are
consing an identical value onto an existing `C a`

, and if so, simply
returns the `C a`

unchanged. If we always use `(::=)`

instead of
directly using `Cons`

, it ensures that `River`

values are always
*normalized*—that is, for every eventually constant stream, we
always use the canonical representative where the element immediately
preciding the constant tail is not equal to it.

This, in turn, *technically* makes it impossible to write functions
which do not respect the equivalence `C a === Cons a (C a)`

, simply
because they will only ever be given canonical rivers as input.
However, as we will see when we discuss folds, it is still possible to
write “bad” functions, *i.e.* functions that are semantically
questionable as functions on eventually constant streams—it would
just mean we cannot directly observe them behaving badly.

The big downside of using this formulation is that the `Eq`

constraint
infects absolutely everything—we even end up with `Eq`

constraints
in places where we would not expect them (for example, on `head :: River a -> a`

), because the pattern synonym incurs an `Eq`

constraint
anywhere we use it, regardless of whether we are using it to construct
or destruct `River`

values. As you can see from the definition above,
we only do an equality check when using `(::=)`

to construct a
`River`

, not when using it to pattern-match, but there is no way to
give the pattern synonym different types in the two directions.Of course, we could make it a unidirectional pattern synonym and just make a differently named smart constructor, but that seems somewhat ugly, as we would have to remember which to use in which situation.

So, because this normalizing variant does not really go far enough in
removing our burden of proof, and has some big downsides in the form
of leaking `Eq`

constraints everywhere, I have chosen to stick with
the simpler `(:::)`

in this post. But I am still a bit unsure about this
choice; in fact, I went back and forth two times while writing.

We can at least provide a `normalize`

function, which we can use when
we want to ensure normalization:

```
normalize :: Eq a => River a -> River a
C a) = C a
normalize (::= as) = a ::= as normalize (a
```

With the preliminary definitions out of the way, we can now build up a
library of standard functions and instances for working with `River a`

values. To start, we can write an `Eq`

instance as follows:

```
instance Eq a => Eq (River a) where
C a == C b = a == b
::: as) == (b ::: bs) = a == b && as == bs (a
```

Notice that we only need two cases, not four: if we compare two values
whose finite prefixes are different lengths, the shorter one will
automatically expand (via matching on `(:::)`

) to the length of the
longer.

We already derived a `Functor`

instance; we can also define a “zippy”
`Applicative`

instance like so:

```
repeat :: a -> River a
repeat = C
instance Applicative River where
pure = repeat
C f <*> C x = C (f x)
::: fs) <*> (x ::: xs) = f x ::: (fs <*> xs)
(f
zipWith :: (a -> b -> c) -> River a -> River b -> River c
zipWith = liftA2
```

We can write safe `head`

, `tail`

, and index functions:

```
head :: River a -> a
head (a ::: _) = a
tail :: River a -> River a
tail (_ ::: as) = as
infixl 9 !!
(!!) :: River a -> Int -> a
C a !! _ = a
::: _) !! 0 = a
(a ::: as) !! n = as !! (n - 1) (_
```

We can also write `take`

and `drop`

variants. Note that `take`

returns a finite prefix of a `River`

, which is a list, not another
`River`

. The special case for `drop _ (C a)`

is not strictly
necessary, but makes it more efficient.

```
take :: Int -> River a -> [a]
take n _ | n <= 0 = []
take n (a ::: as) = a : take (n - 1) as
drop :: Int -> River a -> River a
drop n r | n <= 0 = r
drop _ (C a) = C a
drop n (_ ::: as) = drop (n - 1) as
```

There are many other such functions we could implement (*e.g.* `span`

,
`dropWhile`

, `tails`

…); if I eventually put this on Hackage I would
be sure to have a much more thorough selection of functions. Which
functions would you want to see?

`River`

How do we fold over a `River a`

? The `Foldable`

type class requires us
to define either `foldMap`

or `foldr`

; let’s think about `foldMap`

,
which would have type

`foldMap :: Monoid m => (a -> m) -> River a -> m`

However, this doesn’t really make sense. For example, suppose we have
a `River Int`

; if we had `foldMap`

with the above type, we could use
`foldMap Sum`

to turn our `River Int`

into a `Sum Int`

. But what is
the sum of an infinite stream of `Int`

? Unless the eventually
repeating part is `C 0`

, this is not well-defined. If we simply write
a function to add up all the `Int`

values in a `River`

, including
(once) the value contained in the final `C`

, this would be a good
example of a semantically “bad” function: it does not respect the law
`C a === a ::: C a`

. If we ensure `River`

values are always
normalized, we would not be able to directly observe anything amiss,
but the function still seems suspect.

Thinking about the law `C a === a ::: C a`

again is the key.
Supposing `foldMap f (C a) = f a`

(since it’s unclear what else it
could possibly do), applying `foldMap`

to both sides of the law we
obtain `f a == f a <> f a`

, that is, the combining operation must be
*idempotent*. This makes sense: with an idempotent operation,
continuing to apply the operation to the infinite constant tail will
not change the answer, so we can simply stop once we reach the `C`

.

We can create a subclass of `Semigroup`

to represent *idempotent*
semigroups, that is, semigroups for which `a <> a = a`

. There are
several idempotent semigroups in `base`

; we list a few below. Note
that since rivers are never empty, we can get away with just a
semigroup and not a monoid, since we do not need an identity value
onto which to map an empty structure.

```
class Semigroup m => Idempotent m
-- No methods, since Idempotent represents adding only a law,
-- namely, ∀ a. a <> a == a
-- Exercise for the reader: convince yourself that these are all
-- idempotent
instance Idempotent All
instance Idempotent Any
instance Idempotent Ordering
instance Ord a => Idempotent (Max a)
instance Ord a => Idempotent (Min a)
```

Now, although we cannot make a `Foldable`

instance, we can write our own
variant of `foldMap`

which requires an idempotent semigroup instead of
a monoid:

```
foldMap :: Idempotent m => (a -> m) -> River a -> m
foldMap f (C a) = f a
foldMap f (a ::: as) = f a <> foldMap f as
fold :: Idempotent m => River m -> m
= foldMap id fold
```

We can then instantiate it at some of the semigroups listed above to
get some useful folds. These are all guaranteed to terminate and
yield a sensible answer on any `River`

.

```
and :: River Bool -> Bool
and = getAll . foldMap All
or :: River Bool -> Bool
or = getAny . foldMap Any
all :: (a -> Bool) -> River a -> Bool
all f = and . fmap f
any :: (a -> Bool) -> River a -> Bool
any f = or . fmap f
maximum :: Ord a => River a -> a
maximum = getMax . foldMap Max
minimum :: Ord a => River a -> a
minimum = getMin . foldMap Min
lexicographic :: Ord a => River a -> River a -> Ordering
= fold $ zipWith compare xs ys lexicographic xs ys
```

We could make an `instance Ord a => Ord (River a)`

with `compare = lexicographic`

; however, in the next section I want to make a
different `Ord`

instance for a specific instantiation of `River`

.

Briefly, here’s the particular application I have in mind:
infinite-precision two’s complement arithmetic, *i.e.* \(2\)-adic
numbers. Chris Smith also wrote about \(2\)-adic numbers
recently;
however, unlike Chris, I am not interested in \(2\)-adic numbers in
general, but only specifically those \(2\)-adic numbers which represent
an embedded copy of \(\mathbb{Z}\). These are precisely the eventually
constant ones: nonnegative integers are represented in binary as
usual, with an infinite tail of \(0\) bits, and negative integers are
represented with an infinite tail of \(1\) bits. For example, \(-1\) is
represented as an infinite string of all \(1\)’s. The amazing thing
about this representation (and the reason it is commonly used in
hardware) is that the usual addition and multiplication algorithms
continue to work without needing special cases to handle negative
integers. If you’ve never seen how this works, you should definitely
read
about it.

```
data Bit = O | I deriving (Eq, Ord, Enum)
type Bits = River Bit
```

First, some functions to convert to and from integers. We only need
special cases for \(0\) and \(-1\), and beyond that it is just the usual
business with `mod`

and `div`

to peel off one bit at a time, or
multiplying by two and adding to build up one bit at a time. (I am a big fan of `LambdaCase`

.)

```
toBits :: Integer -> Bits
= \case
toBits 0 -> C O
-1 -> C I
-> toEnum (fromIntegral (n `mod` 2)) ::: toBits (n `div` 2)
n
fromBits :: Bits -> Integer
= \case
fromBits C O -> 0
C I -> -1
::: bs -> 2 * fromBits bs + fromIntegral (fromEnum b) b
```

For testing, we can also make a `Show`

instance. When it comes to
showing the infinite constant tail, I chose to repeat the bit 3 times
and then show an ellipsis; this is not really necessary but somehow
helps my brain more easily see whether it is an infinite tail of zeros
or ones.

```
instance Show Bits where
show = reverse . go
where
C b) = replicate 3 (showBit b) ++ "..."
go (::: bs) = showBit b : go bs
go (b
= ("01" P.!!) . fromEnum showBit
```

Let’s try it out:

```
ghci> toBits 26
...00011010
ghci> toBits (-30)
...11100010
ghci> fromBits (toBits (-30))
-30
ghci> quickCheck $ \x -> fromBits (toBits x) == x
+++ OK, passed 100 tests.
```

Let’s implement some arithmetic. First, incrementing. It is standard
except for a special case for `C I`

(without which, incrementing `C I`

would diverge). Notice that we use `(::=)`

instead of `(:::)`

, which
ensures our `Bits`

values remain normalized.

```
inc :: Bits -> Bits
= \case
inc C I -> C O
O ::= bs -> I ::= bs
I ::= bs -> O ::= inc bs
```

`dec`

is similar, just the opposite:

```
dec :: Bits -> Bits
= \case
dec C O -> C I
I ::= bs -> O ::= bs
O ::= bs -> I ::= dec bs
```

Then we can write `inv`

to invert all bits, and `neg`

as the
composition of `inc`

and `inv`

.

```
inv :: Bits -> Bits
= fmap $ \case { O -> I; I -> O }
inv
neg :: Bits -> Bits
= inc . inv neg
```

Trying it out:

```
λ> toBits 3
...00011
λ> neg it
...11101
λ> inc it
...1110
λ> inc it
...111
λ> inc it
...000
λ> inc it
...0001
λ> dec it
...000
λ> dec it
...111
```

Finally, addition, multiplication, and `Ord`

and `Num`

instances:

```
add :: Bits -> Bits -> Bits
= \cases
add C O) y -> y
(C O) -> x
x (C I) (C I) -> O ::= C I
(I ::= xs) (I ::= ys) -> O ::= inc (add xs ys)
(::= xs) (y ::= ys) -> (x .|. y) ::= add xs ys
(x where
I .|. _ = I
.|. y = y
_
mul :: Bits -> Bits -> Bits
= \cases
mul C O) _ -> C O
(C O) -> C O
_ (C I) y -> neg y
(C I) -> neg x
x (O ::= xs) ys -> O ::= mul xs ys
(I ::= xs) ys -> add ys (O ::= mul xs ys)
(
instance Ord Bits where
-- It's a bit mind-boggling that this works
compare (C x) (C y) = compare y x
compare (x ::= xs) (y ::= ys) = compare xs ys <> compare x y
instance Num Bits where
fromInteger = toBits
negate = neg
+) = add
(*) = mul
(abs = toBits . abs . fromBits
signum = toBits . signum . fromBits
```

```
λ> quickCheck $ withMaxSuccess 1000 $ \x y -> fromBits (mul (toBits x) (toBits y)) == x * y
+++ OK, passed 1000 tests.
λ> quickCheck $ \x y -> compare (toBits x) (toBits y) == compare x y
+++ OK, passed 100 tests.
```

Just for fun, let’s implement the Collatz map:

```
collatz :: Bits -> Bits
O ::= bs) = bs
collatz (@(I ::= _) = 3*bs + 1 collatz bs
```

```
λ> P.take 20 $ map fromBits (iterate collatz (toBits (-13)))
[-13,-38,-19,-56,-28,-14,-7,-20,-10,-5,-14,-7,-20,-10,-5,-14,-7,-20,-10,-5]
λ> P.take 20 $ map fromBits (iterate collatz (toBits 7))
[7,22,11,34,17,52,26,13,40,20,10,5,16,8,4,2,1,4,2,1]
```

Is

`(:::)`

or`(::=)`

the better default? It’s tempting to just say “provide both and let the user decide”. I don’t disagree with that; however, the question is which one we use to implement various basic functions such as`map`

/`fmap`

. For example, if we use`(:::)`

, we can make a`Functor`

instance, but values may not be normalized after mapping.Can we generalize from eventually constant to eventually

*periodic*? That is, instead of repeating the same value forever, we cycle through a repeating period of some finite length. I think this is possible, but it would make the implementation more complex, and I don’t know the right way to generalize`foldMap`

. (We could insist that it only works for*commutative*idempotent semigroups, but in that case what’s the point of having a*sequence*of values rather than just a set?)

Happy to hear any comments or suggestions!

In a previous
post
I challenged you to solve Factor-Full
Tree. In this
problem, we are given an unlabelled rooted tree, and asked to create a *divisor
labelling*. That is, we must label the vertices with positive
integers in such a way that \(v\) is an ancestor of \(u\) if and only if
\(v\)’s label evenly divides \(u\)’s label.

For example, here is a tree with a divisor labelling:

Divisor labelling of a tree

The interesting point (though irrelevant to solving the problem) is
that this is a method for encoding a tree as a set of integers:
because \(v\) is an ancestor of \(u\) *if and only if* \(v\)’s label divides
\(u\)’s, all the information about the tree’s structure is fully
contained in the set of labels. For example, if we simply write
down the set \(\{1, 5, 6, 7, 12, 14, 21, 49, 63\}\), it is possible to
fully reconstruct the above tree from this set.Note that we
consider trees equivalent up to reordering of siblings, that is, each
node has a *bag*, not a *list*, of children.

This is not a
particularly *efficient* way to encode a tree, but it is certainly
interesting!

First, some basic setup.See here for the `Scanner`

abstraction, and
here
for the basics of how I organize solutions.

The first line of
input specifies the number of nodes \(N\), and after that there are
\(N-1\) lines, each specifying a single undirected edge.

```
import Control.Category ((>>>))
import Data.Bifunctor (second)
import Data.Map (Map, (!?))
import qualified Data.Map as M
import Data.Tuple (swap)
= C.interact $ runScanner tc >>> solve >>> format
main
data TC = TC { n :: Int, edges :: [Edge] }
deriving (Eq, Show)
tc :: Scanner TC
= do
tc <- int
n <- (n - 1) >< pair int int
edges return TC{..}
format :: [Integer] -> ByteString
= map showB >>> C.unwords format
```

We are guaranteed that the edges describe a tree; next we will actually build a tree data structure from the input.

There are many similar problems which specify a tree structure by giving a list of edges, so it’s worthwhile trying to write some generic code to transform such an input into an actual tree. In an imperative language we would do this by building a map from each node to its neighbors, then doing a DFS to orient the tree. Our Haskell code will be similar, except building the map and doing a DFS will both be one-liners!

First, a function to turn a list of undirected edges into a `Map`

associating each vertex to all its neighbors. It’s convenient to
decompose this into a function to turn a list of *directed* edges into
a `Map`

, and a function to duplicate and swap each pair. We won’t
need `dirEdgesToMap`

for this problem, but we can certainly imagine
wanting it elsewhere.

```
edgesToMap :: Ord a => [(a, a)] -> Map a [a]
= concatMap (\p -> [p, swap p]) >>> dirEdgesToMap
edgesToMap
dirEdgesToMap :: Ord a => [(a, a)] -> Map a [a]
= map (second (: [])) >>> M.fromListWith (++) dirEdgesToMap
```

Next, we can turn such a neighbor `Map`

into a tree. Rather than
returning a literal `Tree`

data structure, it’s convenient to
incorporate a tree fold: that is, given a function `a -> [b] -> b`

, a neighbor
map, and a root node, we fold over the whole tree and return the
resulting `b`

value. (Of course, if we want an actual `Tree`

we can use
`mapToTree Node`

.) We can also compose these into a single function `edgesToTree`

.

```
mapToTree :: Ord a => (a -> [b] -> b) -> Map a [a] -> a -> b
= dfs root root
mapToTree nd m root where
= nd root (maybe [] (map (dfs root) . filter (/= parent)) (m !? root))
dfs parent root
edgesToTree :: Ord a => (a -> [b] -> b) -> [(a, a)] -> a -> b
= mapToTree nd . edgesToMap edgesToTree nd
```

So how do we create a divisor labelling for a given tree? Clearly, we
might as well choose the root to have label \(1\), and every time we
descend from a parent to a child, we must multiply by some integer,
which might as well be a prime. Of course, we need to multiply by a
*different* prime for each sibling. We might at first imagine simply
multiplying by 2 for each (arbitrarily chosen) leftmost child, 3 for
each second child, 5 for each third child, and so on, but this does
not work—the second child of the first child ends up with the same
label as the first child of the second child, and so on.

Each node \(u\)’s label is some prime \(p\) times its parent’s label; call
\(p\) the *factor* of node \(u\). It is OK for one child of \(u\) to also
have factor \(p\), but the other children must get different factors.
To be safe, we can give each additional child a new *globally unique*
prime factor. This is not always necessary—in some cases it can be
OK to reuse a factor if it does not lead to identically numbered
nodes—but it is certainly sufficient. As an example, below is a
divisor labelling of the example tree from before, via this scheme.
Each edge is labelled with the factor of its child.

Divisor labelling of a tree with consecutive primes

Notice how we use \(2\) for the first child of the root, and \(3\) for the next child. \(3\)’s first child can also use a factor of \(3\), yielding a label of \(3^2 = 9\). \(3\)’s next child uses a new, globally unique prime \(5\), and its third child uses \(7\); the final child of \(1\) uses the next available prime, \(11\).

We can code this up via a simple stateful traversal of the tree. (For
`primes`

, see this
post.)
It’s a bit fiddly since we have to switch to the next prime *between*
consecutive children, but not *after* the last child.

```
primes :: [Integer]
= 2 : sieve primes [3 ..]
primes where
: ps) xs =
sieve (p let (h, t) = span (< p * p) xs
in h ++ sieve ps (filter ((/= 0) . (`mod` p)) t)
curPrime :: State [Integer] Integer
= gets head
curPrime
nextPrime :: State [Integer] ()
= modify tail
nextPrime
labelTree :: Tree a -> Tree (Integer, a)
= flip evalState primes . go 1
labelTree where
go :: Integer -> Tree a -> State [Integer] (Tree (Integer, a))
Node a ts) = Node (x, a) <$> labelChildren x ts
go x (
labelChildren :: Integer -> [Tree a] -> State [Integer] [Tree (Integer, a)]
= pure []
labelChildren _ [] : ts) = do
labelChildren x (t <- curPrime
p <- go (x * p) t
t' case ts of
-> pure [t']
[] -> do
_
nextPrime:) <$> labelChildren x ts (t'
```

There is a bit of additional glue code we need get the parsed tree
from the input, apply `labelTree`

, and then print out the node
labels in order. However, I’m not going to bother showing it,
because—this solution is not accepted! It fails with a WA (Wrong
Answer) verdict. What gives?

The key is one of the last sentences in the problem statement, which I haven’t mentioned so far: all the labels in our output must be at most \(10^{18}\). Why is this a problem? Multiplying by primes over and over again, it’s not hard to get rather large numbers. For example, consider the tree below:

Tree for which our naïve scheme generates labels that are too large

Under our scheme, the root gets label \(1\), and the children of the root get consecutive primes \(2, 3, 5, \dots, 29\). Then the nodes in the long chain hanging off the last sibling get labels \(29^2, 29^3, \dots, 29^{13}\), and \(29^{13}\) is too big—in fact, it is approximately \(10^{19}\). And this tree has only 23 nodes; in general the input can have up to 60.

Of course, \(29\) was a poor choice of factor for such a long chain—we should have instead labelled the long chain with powers of, say, 2. Notice that if we have a “tree” consisting of a single long chain of 60 nodes (and you can bet this is one of the secret test inputs!), we just barely get by labelling it with powers of two from \(2^0\) up to \(2^{59}\): in fact \(2^{59} < 10^{18} < 2^{60}\). So in general, we want to find a way to label long chains with small primes, and reserve larger primes for shorter chains.

One obvious approach is to simply sort the children at each node by decreasing height, before traversing the tree to assign prime factors. This handles the above example correctly, since the long chain would be sorted to the front and assigned the factor 2. However, this does not work in general! It can still fail to assign the smallest primes to the longest chains. As a simple example, consider this tree, in which the children of every node are already sorted by decreasing height from left to right:

Tree for which sorting by height first does not work

The straightforward traversal algorithm indeed assigns powers of 2 to the left spine of the tree, but it then assigns 3, 5, 7, and so on to all the tiny spurs hanging off it. So by the time we get to other long chain hanging off the root, it is assigned powers of \(43\), which are too big. In fact, we want to assign powers of 2 to the left spine, powers of 3 to the chain on the right, and then use the rest of the primes for all the short spurs. But this sort of “non-local” labelling means we can’t assign primes via a tree traversal.

To drive this point home, here’s another example tree. This one is small enough that it probably doesn’t matter too much how we label it, but it’s worth thinking about how to label the longest chains with the smallest primes. I’ve drawn it in a “left-leaning” style to further emphasize the different chains that are involved.

Tree with chains of various lengths

In fact, we want to assign the factor 2 to the long chain on the left; then the factor 3 to the second-longest chain, in the fourth column; then 5 to the length-6 chain in the second column; 7 to the length-3 chain all the way on the right; and finally 11 to the smallest chain, in column 3.

In general, then, we want a way to *decompose* an arbitrary tree into
chains, where we repeatedly identify the longest chain, remove it from
consideration, and then identify the longest chain from the remaining
nodes, and so on. Once we have decomposed a tree into chains, it will
be a relatively simple matter to sort the chains by length and assign
consecutive prime factors.

This decomposition occasionally comes in handy (for example, see Floating Formation), and belongs to a larger family of important tree decomposition techniques such as heavy-light decomposition. Next time, I’ll demonstrate how to implement such tree decompositions in Haskell!

Recently, Dani Rybe wrote this really cool blog
post
(in turn based on this old post by Samuel
Gélineau)
about encoding truly *unordered* n-tuples in Haskell. This is
something I thought about a long time ago in my work on
combinatorial species, but I never came up with a way to represent
them. Samuel and Dani’s solution is wonderful and clever and totally
impractical, and I love it.

I won’t go into more detail than that; I’ll let you go read it if you’re interested. This blog post exists solely to respond to Dani’s statement towards the end of her post:

I’m not sure how to, for example, write a function that multiplies the inputs.

Challenge accepted!

```
primes :: [Int]
= 2 : sieve primes [3 ..]
primes where
: ps) xs =
sieve (p let (h, t) = span (< p * p) xs
in h ++ sieve ps (filter ((/= 0) . (`mod` p)) t)
mul :: [Int] -> Int
= unfuck mulU
mul where
mulU :: U n Int -> Int
= ufold 1 id (< 0) \(US neg nonNeg) ->
mulU * mulPos primes (abs <$> neg) * (-1) ^ ulen neg
mulNonNeg nonNeg
mulNonNeg :: U n Int -> Int
= ufold 1 id (== 0) \(US zero pos) ->
mulNonNeg if ulen zero > 0 then 0 else mulPos primes pos
mulPos :: [Int] -> U n Int -> Int
= ufold 1 id (== 1) \(US _ pos) -> mulGTOne ps pos
mulPos ps
mulGTOne :: [Int] -> U n Int -> Int
: ps) = ufold 1 id ((== 0) . (`mod` p)) \(US divP nondivP) ->
mulGTOne (p : ps) ((`div` p) <$> divP) * (p ^ ulen divP) * mulGTOne ps nondivP mulPos (p
```

Since every integer has a unique prime factorization, at each step we split the remaining numbers into those divisible by \(p\) and those not divisible by \(p\). For the ones that are, we divide out \(p\) from all of them, multiply by the appropriate power of \(p\), and recurse on what’s left; for those that are not, we move on to trying the next prime.

Dani also speculates about `ubind :: U n (U m a) -> U (n :*: m) a`

. I
believe in my heart this should be possible to implement, but after
playing with it a bit, I concluded it would require an astounding feat
of type-fu.

PS I’m working on getting comments set up here on my new blog… hopefully coming soon!

In a previous post I challenged you to solve Product Divisors. In this problem, we are given a sequence of positive integers \(a_1, \dots, a_n\), and we are asked to compute the total number of divisors of their product. For example, if we are given the numbers \(4, 2, 3\), then the answer should be \(8\), since \(4 \times 2 \times 3 = 24\) has the \(8\) distinct divisors \(1, 2, 3, 4, 6, 8, 12, 24\).

In general, if \(a\) has the prime factorization \(a = p_1^{\alpha_1} p_2^{\alpha_2} \cdots p_k^{\alpha_k}\) (where the \(p_i\) are all distinct primes), then the number of divisors of \(a\) is

\[(\alpha_1 + 1)(\alpha_2 + 1) \cdots (\alpha_k + 1),\]

since we can independently choose how many powers of each prime to include. There are \(\alpha_i + 1\) choices for \(p_i\) since we can choose anything from \(p_i^0\) up to \(p_i^{\alpha_i}\), inclusive.

So at a fundamental level, the solution is clear: factor each \(a_i\),
count up the number of copies of each prime in their product, then do
something like `map (+1) >>> product`

. We are also told the answer
should be given mod \(10^9 + 7\), so we can use aUsing `Int`

instead of `Integer`

here is OK as long as we are sure to be running
on a 64-bit system; multiplying two `Int`

values up to \(10^9 + 7\)
yields a result that still fits within a 64-bit signed `Int`

.
Otherwise (*e.g.* on Codeforces) we would have to use `Integer`

.

`newtype`

with a
custom `Num`

instance:

```
p :: Int
= 10^9 + 7
p
newtype M = M { unM :: Int } deriving (Eq, Ord)
instance Show M where show = show . unM
instance Num M where
fromInteger = M . (`mod` p) . fromInteger
M x + M y = M ((x + y) `mod` p)
M x - M y = M ((x - y) `mod` p)
M x * M y = M ((x * y) `mod` p)
```

Of course, I would not be writing about this problem if it were that
easy! If we try implementing the above solution idea in a
straightforward way—for example, if we take the simple factoring code from this blog
post
and then do something like `map factor >>> M.unionsWith (+) >>> M.elems >>> map (+1) >>> product`

, we get the dreaded Time Limit Exceeded.

Why doesn’t this work? I haven’t mentioned how many integers might be in the input: in fact, we might be given as many as one million (\(10^6\))! We need to be able to factor each number very quickly if we’re going to finish within the one second time limit. Factoring each number from scratch by trial division is simply too slow.

While more sophisticated methods are needed to factor a *single*
number more quickly than trial division, there is a standard technique
we can use to speed things up when we need to factor *many* numbers.
We can use a *sieve* to precompute a lookup table, which we can then
use to factor numbers very quickly.

In particular, we will compute a table \(\mathit{smallest}\) such that
\(\mathit{smallest}[i]\) will store the *smallest prime factor* of \(i\).
Given this table, to factor a positive integer \(i\), we simply look up
\(\mathit{smallest}[i] = p\), add it to the prime factorization, then
recurse on \(i/p\); the base case is when \(i = 1\).

How do we compute \(\mathit{smallest}\)? The basic idea is to create an
array of size \(n\), initializing it with \(\mathit{smallest}[k] = k\). For each \(k\) from \(2\) up to \(n\),We could optimize this even
further via the approach in this blog
post, which takes \(O(n)\)
rather than \(O(n \lg n)\) time, but it would complicate our Haskell
quite a bit and it’s not needed for solving this problem.

if
\(\mathit{smallest}[k]\) is still equal to \(k\), then \(k\) must be prime;
iterate through multiples of \(k\) (starting with \(k^2\), since any
smaller multiple of \(k\) is already divisible by a smaller prime) and
set each \(\mathit{smallest}[ki]\) to the minimum of \(k\) and whatever
value it had before.

This is one of those cases where for efficiency’s sake, we actually
want to use an honest-to-goodness mutable array. Immutable arrays are
not a good fit for sieving, and using something like a `Map`

would
introduce a lot of overhead that we would rather avoid. However, we
only need the table to be mutable while we are computing it; after
that, it should just be an immutable lookup table. This is a great fit
for an `STUArray`

:Note that as of this writing, the version of the
`array`

library installed in the Kattis environment does not have
`modifyArray'`

, so we actually have to do `readArray`

followed by
`writeArray`

.

```
= 1000000
maxN
smallest :: UArray Int Int
= runSTUArray $ do
smallest <- newListArray (2,maxN) [2 ..]
a 2 .. maxN] $ \k -> do
forM_ [<- readArray a k
k' == k') $ do
when (k *k, k*(k+1) .. maxN] $ \n ->
forM_ [kmin k)
modifyArray' a n (return a
```

Haskell, the world’s finest imperative programming language!

We can now write a new `factor`

function that works by repeatedly
looking up the smallest prime factor:

```
factor :: Int -> Map Int Int
= \case
factor 1 -> M.empty
-> M.insertWith (+) p 1 (factor (n `div` p))
n where
= smallest!n p
```

And now we can just do `map factor >>> M.unionsWith (+) >>> M.elems >>> map (+1) >>> product`

as before, but since our `factor`

is so much faster this time, it
should…

What’s that? Still TLE? Sigh.

Unfortunately, creating a bunch of `Map`

values and then doing
`unionsWith`

one million times still introduces way too much overhead.
For many problems working with `Map`

(which is impressively fast) is
good enough, but not in this case. Instead of returning a `Map`

from
each call to `factor`

and then later combining them, we can write a version of
`factor`

that directly increments counters for each prime in a
mutable array:

```
factor :: STUArray s Int Int -> Int -> ST s ()
= go n
factor counts n where
1 = return ()
go = do
go n let p = smallest!n
+1)
modifyArray' counts p (`div` p) go (n
```

Then we have the following top-level solution, which is finally fast enough:

```
main :: IO ()
= C.interact $ runScanner (numberOf int) >>> solve >>> showB
main
solve :: [Int] -> M
= counts >>> elems >>> map ((+1) >>> M) >>> product
solve
counts :: [Int] -> UArray Int Int
= runSTUArray $ do
counts ns <- newArray (2,maxN) 0
cs
forM_ ns (factor cs)return cs
```

This solution runs in just over 0.4s for me. Considering that this is
only about 4x slower than the fastest solution (0.09s, in C++), I’m
pretty happy with it! We did have to sacrifice a bit of elegance for
speed, especially with the `factor`

and `counts`

functions instead of
`M.unionsWith`

, but in the end it’s not too bad.

I thought we might be able to make this even faster by using a strict
fold over the `counts`

array instead of converting to a list with
`elems`

and then doing a `map`

and a `product`

, but (1) there is no
generic fold operation on `UArray`

, and (2) I trust that GHC is
already doing a pretty good job optimizing this via list fusion.

Next time I’ll write about my solution to the other challenge problem, Factor-Full Tree. Until then, give it a try!