Connecting cities with Union-Find

September 15, 2020

Let’s assume we have a grid $G$ of size $r \times c$ . We have $n$ cities placed on blocks in $G$ . We task players Alice & Bob with placing tiny train tracks, one on a block, between these cities in order to connect them to each other. How can we determine which cities are connected?

After Alice and Bob placed their tracks

Introducing Union-Find

A Union-Find data structure wraps around a list of items and provides us with two actions:

// Declare that a and b belong to the same set
union<T>(a: T, b: T): void;

// Does a and b belong to the same set?
find<T>(a: T, b: T): boolean;

Note: For a more visual representation of how the algorithm works, see Techie’s Delight write-up.

It’s important to note that Union-Find asks and answers the question “Does A and B belong to the same set, yes or no?” and not how they are connected to each other. This reduces the algorithm’s internal complexity akin to a simple lookup.

Let’s write some code

Note: This is an abridged version of the complete implementation. For the sake of brevity we exclude error checking, path compression and tree balancing. These are interesting topics to read up on, and I’d highly recommend you do so should you choose to implement Union-Find.

| Initial Setup

const DEFAULT_SERIALIZER = JSON.stringify;

class UnionFind<T> {
  private roots: number[] = [];
  // Since we specify roots with array indices, we keep a map of the object's value to its index
  private itemToIndexMap: { [key: string]: number } = {};

  constructor(private serializer: (item: T) => string = DEFAULT_SERIALIZER) {}
}

| Registering an Item

Since we use a serializer and itemToIndexMap to keep track of the items in our sets, we need to register a new item before using it for the first time.

register(item: T): void {
  const newIndex = this.roots.length;
  const serialized = this.serializer(item);
  this.itemToIndexMap[serialized] = newIndex;
  // Initially, an item's root is itself
  this.roots.push(newIndex);
}

| Querying whether an item as been registered

isItemRegistered(item: T): boolean {
  const serialized = this.serializer(item);
  return this.itemToIndexMap[serialized] !== undefined;
}

| Finding the root of an item

root(item: T): number {
  // First we convert the item to its index representation
  const indexRepresentation = this.itemToIndexMap[this.serializer(item)];

  // The search for the root parent begins
  let parent = this.roots[indexRepresentation];
  
  // We stop when we find the item whose parent is itself, i.e. the root
  while (this.roots[parent] !== parent) {
    parent = this.roots[parent];
  }
  return parent;
}

| Union

union(a: T, b: T): void {
  const rootA = this.root(a);
  const rootB = this.root(b);

  // Set the root of A to the root of B
  this.roots[rootA] = rootB;
}

| Find

find(a: T, b: T): boolean {
  // Does A and B have the same roots?
  return this.root(a) === this.root(b);
}

Connecting Cities

Visual Representation

In this example, we’re building this grid:

After Alice and Bob placed their tracks

Setup

interface Coordinate {
  r: number;
  c: number;
}

const aliceTracks: Coordinate[] = [
  { r: 1, c: 3 },
  { r: 1, c: 1 },
  { r: 3, c: 4 },
  { r: 4, c: 4 },
]

const bobTracks: Coordinate[] = [
  { r: 1, c: 2 },
  { r: 2, c: 1 },
  { r: 2, c: 4 },
  { r: 4, c: 5 },
  { r: 4, c: 6 },
];

const cities: Coordinate[] = [
  { r: 1, c: 4 },
  { r: 3, c: 1 },
  { r: 4, c: 7 },
]

// The default serializer is sufficient
const unionFind = new UnionFind<Coordinate>();

Registering our items

Let’s assume we don’t care whether a track was placed by Alice or Bob.

// Register all cities and tracks
[...cities, ...aliceTracks, ...bobTracks]
  .forEach(unionFind.register)

Declaring their unions

For a given coordinate, we need to find all of its neighbors. For simplicity, we ignore diagonal neighbors.

const allAdjacentCoordinates = ({ r, c }: Coordinate): Coordinate[] => {
  return [
    { r: r - 1, c },
    { r: r + 1, c },
    { r, c: c - 1 },
    { r, c: c + 1 },
  ];
};


[...cities, ...aliceTracks, ...bobTracks].forEach(coordinate => {
  allAdjacentCoordinates(coordinate)
    .filter(unionFind.isItemRegistered)
    .forEach(neighbor => {
      unionFind.union(coordinate, neighbor);
    });
});

Asking all the right questions

To see whether two cities are connected to each other, we simply ask whether they belong to the same set.

if (unionFind.find(cities[0], cities[1])) {
  console.log("Hooray, they're connected. 🎉");
}

Distinguishing between Alice and Bob

What if we wanted to distinguish between Alice & Bob, and we wanted to know whether Alice has connected two cities together?

We keep separate UnionFind instances for Alice & Bob.

const unionFindAlice = new UnionFind<Coordinate>();
const unionFindBob = new UnionFind<Coordinate>();

// Register all cities and tracks for Alice
[...cities, ...aliceTracks]
  .forEach(unionFindAlice.register)

// Register all cities and tracks for Bob
[...cities, ...bobTracks]
  .forEach(unionFindBob.register)

// Declare unions for Alice
[...cities, ...aliceTracks].forEach(coordinate => {
  allAdjacentCoordinates(coordinate)
    .filter(unionFindAlice.isItemRegistered)
    .forEach(neighbor => {
      unionFindAlice.union(coordinate, neighbor);
    });
});

// Declare unions for Bob
[...cities, ...bobTracks].forEach(coordinate => {
  allAdjacentCoordinates(coordinate)
    .filter(unionFindBob.isItemRegistered)
    .forEach(neighbor => {
      unionFindBob.union(coordinate, neighbor);
    });
});


if (unionFindAlice.find(cities[0], cities[1])) {
  console.log("Hooray, Alice has connected these cities. 🎉");
}