Number of islands II: union-find on a growing grid in Python

Why this matters

A grid is a table of cells. Each cell is either water or land. A group of land cells that touch side to side is an island. Two cells touch when they share an edge. Diagonal cells do not count as touching.

An earlier lesson in this module counted islands on a fixed grid. There, the whole map was known in advance, and a flood fill (a walk that visits every connected land cell) found each island once.

This lesson changes the question. The grid starts all water. Land arrives one cell at a time, and after each arrival you need the current island count. You could run a full flood fill after every arrival. That works, but each run looks at every cell on the map, so the total work grows with both the number of arrivals and the size of the grid.

Union-find is a better tool here. It keeps track of which cells belong to the same island. When a new cell arrives, it only needs to look at the four cells next to it and merge islands where needed. The rest of the map is never visited.

Terms used in this lesson

A grid with rows rows and cols columns has cells labelled by their row and column. Counting starts at 0, so the top-left cell is (0, 0). The cell in row r and column c is written (r, c).

A neighbour of a cell is one of the four cells directly above, below, left or right of it. Cells at the edge of the grid have fewer neighbours.

An island is a set of land cells where you can walk from any cell to any other by stepping only between neighbours.

Union-find is a data structure that keeps groups of items. It supports two actions. Find asks which group an item belongs to. Union joins two groups into one.

A parent list is the way union-find stores groups. Each item has a slot in a Python list. The slot holds the item's parent. Following parents from any item leads to a root, an item whose parent is itself. Two items are in the same group when they reach the same root.

A flat index is one number that stands for a cell. For a grid with cols columns, the flat index of (r, c) is r * cols + c. Union-find works on flat indexes, so each cell has one slot in the parent list.

The problem

You are given the grid size and a list of positions. Each position is a cell that becomes land. After each position, report how many islands there are.

Take a 3 by 3 grid and these four positions in order: (0, 0), (0, 1), (1, 2), (2, 1).

Step 2 adds a cell next to the first one, so the island grows and the count stays at 1. Step 3 adds (1, 2). Its neighbours are (0, 2), (2, 2), (1, 1) and the edge of the grid. All of those are water, so this is a new island and the count goes to 2. Step 4 adds (2, 1). Its neighbours (1, 1), (2, 0) and (2, 2) are all water, so again a new island, and the count goes to 3.

The answer for this input is [1, 1, 2, 3]. The list has one number per position.

The idea in three cases

When a new land cell arrives, look at its neighbours. Only land neighbours matter. There are three cases.

Case 1: no land neighbours. The new cell is a new island on its own. The count goes up by 1.

Case 2: land neighbours that all belong to one island. The new cell joins that island. The count does not change.

Case 3: land neighbours that belong to two or more different islands. The new cell joins all of them into one island. The count goes up by 1 for the new cell and down by 1 for each island that is merged into another. The net change is 1 minus the number of islands that merged.

The three cases can be handled with one rule. Add 1 to the count for the new cell. Then, for each land neighbour, union the new cell with that neighbour. Each union that joins two different islands subtracts 1 from the count. A union between two cells that are already in the same island subtracts nothing.

Worked example 1: the 3 by 3 grid

Use the grid from the problem section. Flat indexes for a 3 by 3 grid are: (0, 0) is 0, (0, 1) is 1, (0, 2) is 2, (1, 0) is 3, and so on up to (2, 2), which is 8.

Step 1 adds cell 0. The count goes from 0 to 1. It has no land neighbours, so the count stays at 1.

Step 2 adds cell 1. The count goes to 2. Its neighbour to the left is cell 0, which is land. Union cell 1 with cell 0. They were in different islands, so the count drops back to 1.

Step 3 adds cell 5, which is (1, 2). The count goes to 2. Its neighbours are cell 2 (above), cell 8 (below), cell 4 (left) and nothing on the right. None of them is land, so the count stays at 2.

Step 4 adds cell 7, which is (2, 1). The count goes to 3. Its neighbours are cell 4 (above), cell 6 (left) and cell 8 (right). None is land, so the count stays at 3.

The count after each step is 1, 1, 2, 3. This matches the table.

Worked example 2: a repeated position

The input can list the same cell twice. A cell that is already land is not a new island, so the count must not change.

Take a 3 by 3 grid with positions (0, 0), (0, 0), (0, 1).

Step 1 makes (0, 0) land. The count is 1.

Step 2 is (0, 0) again. The cell is already land, so nothing happens. The count stays at 1.

Step 3 adds (0, 1). The count goes to 2, then the union with the land neighbour (0, 0) subtracts 1. The count is 1.

The answer is [1, 1, 1]. The code checks for this case first and returns early, so repeated positions are never counted twice.

Worked example 3: one cell that joins two islands

This example is harder. It uses a 1 by 5 grid, which is a single row of five cells. The positions are (0, 0), (0, 4), (0, 2), (0, 1), (0, 3).

Step 4 is the interesting one. Column 1 has land on both sides: column 0 and column 2. Those were two separate islands before column 1 arrived. The new cell joins them. The count goes up by 1 to 4, then down by 1 for each union that merges two islands. The first union joins column 1 with column 0, so the count goes from 4 to 3. The second union joins column 1 with column 2, so the count goes from 3 to 2. The result is 2.

Step 5 adds column 3. Its neighbours are column 2, which is in the big island, and column 4, which is a separate island. Joining them leaves one island. The answer is [1, 2, 3, 2, 1].

The union-find part

The code keeps two lists of the same length as the number of cells.

self.parent[x] holds the parent of cell x. When a cell first becomes land, its parent is itself, so it is its own root. Unions change parents so that cells in one island all point, through a chain, to one root.

self.size[x] is used only for roots. It holds the number of cells in the island whose root is x. The code uses it to decide which root goes under the other during a union. The smaller tree goes under the bigger one. This keeps chains short. The lesson on union by size and path halving explains why short chains make find fast.

The three operations are:

• _find(x) walks up the parent links from x until it reaches a cell whose parent is itself. That cell is the root. The loop also makes each cell point one step higher as it goes (parent[x] = parent[parent[x]]). This is called path halving. It keeps later finds short. The basics of find and union are covered in the lesson on union-find and minimum spanning trees.
• _union(a, b) finds the roots of a and b. If they are the same, the two cells are already in one island, and the method does nothing. Otherwise it attaches the smaller tree under the bigger one, adds the sizes, and subtracts 1 from the island count.
• add_land(r, c) sets the cell as land, adds 1 to the count, and calls _union with each land neighbour that is inside the grid.

A leading underscore on a method name, such as _find, is a convention. It tells a reader that the method is for use inside the class.

The code below is the whole program. Save it as islands_ii.py.

Run it with python3 islands_ii.py. The output is:

[1, 1, 2, 3]
[1, 1, 1]
[1, 2, 3, 2, 1]

These are the three worked examples from above. The outputs were produced with Python 3.14.

A few Python points in the code:

• self.parent = list(range(cells)) creates the list [0, 1, 2, ...], so each cell starts as its own root.
• [False] * cells makes a list of cells copies of False.
• 0 <= nr < self.rows is a chained comparison. It means 0 <= nr and nr < self.rows must both be true.
• ra, rb = rb, ra swaps two values in one line.
• [... for r, c in positions] is a list comprehension. It runs the loop body once per item and collects the results into a new list.

Walking through one addition in the code

Take the 1 by 5 example at step 4, when column 1 arrives. Before this step, the three cells 0, 2 and 4 each have their own island. Cell 1 is still water. The state is parent = [0, 1, 2, 3, 4], size = [1, 1, 1, 1, 1], and the island count is 3.

add_land(0, 1) sets land[1] to True and adds 1 to islands. The count is now 4.

The loop checks the four directions in order: up, down, left, right. Up and down are outside the grid, so the if skips them. Left is cell 0. It is land, so _union(1, 0) runs. The roots are 1 and 0. Both trees have size 1, so the swap does not happen. The code attaches root 0 under root 1, and the count drops to 3.

Right is cell 2. It is land, so _union(1, 2) runs. The roots are 1 and 2. The size of root 1 is 2, and the size of root 2 is 1, so no swap happens. The code attaches root 2 under root 1, and the count drops to 2. The function returns 2.

After this step, parent is [1, 1, 1, 3, 4] and size is [1, 3, 1, 1, 1]. Cell 1 is the root of the island that holds cells 0, 1 and 2.

Checking the code against a slow method

A useful habit is to test a fast method against a slow one that is easy to trust. The slow method here runs a flood fill over the whole grid after each addition and counts the islands it finds.

The test below makes 3000 random grids of size up to 6 by 6. It uses a fixed seed, so every run makes the same grids. It adds random positions, including repeats, and compares the union-find answer with the flood-fill answer after every step. Save it as check_islands.py in the same folder as islands_ii.py.

import random
from collections import deque
 
from islands_ii import num_islands_2
 
 
def flood_count(land, rows, cols):
    seen = set()
    count = 0
    for r in range(rows):
        for c in range(cols):
            if land[r][c] and (r, c) not in seen:
                count += 1
                seen.add((r, c))
                queue = deque([(r, c)])
                while queue:
                    cr, cc = queue.popleft()
                    for dr, dc in ((-1, 0), (1, 0), (0, -1), (0, 1)):
                        nr = cr + dr
                        nc = cc + dc
                        if (
                            0 <= nr < rows
                            and 0 <= nc < cols
                            and land[nr][nc]
                            and (nr, nc) not in seen
                        ):
                            seen.add((nr, nc))
                            queue.append((nr, nc))
    return count
 
 
def check(trials=3000, seed=42):
    rng = random.Random(seed)
    for t in range(trials):
        rows = rng.randint(1, 6)
        cols = rng.randint(1, 6)
        k = rng.randint(0, rows * cols * 2)
        positions = [(rng.randrange(rows), rng.randrange(cols)) for _ in range(k)]
        got = num_islands_2(rows, cols, positions)
        land = [[False] * cols for _ in range(rows)]
        for i, (r, c) in enumerate(positions):
            land[r][c] = True
            want = flood_count(land, rows, cols)
            if got[i] != want:
                print(f"MISMATCH trial {t} step {i}: got {got[i]} want {want}")
                return
    print(f"all {trials} random grids match the flood-fill count")
 
 
check()

Run it with python3 check_islands.py. The first three lines of output come from the examples at the bottom of islands_ii.py, because importing the file runs them. The output is:

[1, 1, 2, 3]
[1, 1, 1]
[1, 2, 3, 2, 1]
all 3000 random grids match the flood-fill count

The flood fill is slow for large inputs, so it is only a checker. The union-find version is the one to use.

Common mistakes

Forgetting the bounds check

The bounds check is the line if 0 <= nr < self.rows and 0 <= nc < self.cols:. It stops the code from looking at a neighbour outside the grid. A cell at the bottom edge has no cell below it. Without the check, the code computes a flat index that is past the end of the list.

Here is what happens when the row part of the check is removed and (1, 0) is added to a 2 by 3 grid. The grid has 6 cells, so the list has slots 0 to 5. The cell below (1, 0) is (2, 0), whose flat index is 2 * 3 + 0, which is 6. Index 6 does not exist.

IndexError: list index out of range

The traceback above that line points to the line that reads self.land[neighbour].

Python has one difference from many languages. A negative index does not fail. It counts from the end of the list, so [10, 20, 30][-1] returns 30. A bounds check therefore has to stop negative numbers as well as numbers past the end. The check in this lesson does both.

Removing the column part of the check is worse, because it does not crash. A cell at the right edge has a right neighbour that is outside the grid. Its flat index is the start of the next row. In a 2 by 3 grid, the right neighbour of (0, 2) is flat index 0 * 3 + 3, which is 3. Flat index 3 is cell (1, 0), not a cell outside the grid. So the code joins two cells that are not neighbours.

The version with the column check removed gives this output for the positions (1, 0) and then (0, 2):

[1, 1]

The correct answer is [1, 2]. The two cells are on opposite corners and do not touch. The bug gives the wrong count without any error message, so a test on a few small grids is the only way to catch it.

Counting a repeated position as a new island

If the early return is removed, a position that is already land adds 1 to the count again. The union calls find the cell's neighbours and merge nothing, so the count is simply too high.

With the early return removed, the program gives this output for positions (0, 0), (0, 0), (0, 1):

[1, 2, 2]

The correct answer is [1, 1, 1]. The second position is the same cell, so the island count should not change.

Subtracting for every neighbour instead of every merge

The count goes down only when a union joins two different islands. A common error is to subtract 1 for every land neighbour, without checking whether the neighbour is already in the same island. The class below shows this error. It replaces _union with a method that always subtracts:

class AlwaysSubtract(IslandCounter):
    def _union(self, a, b):
        self.islands -= 1

Take a 2 by 2 grid with positions (0, 0), (0, 1), (1, 0), (1, 1). After the first three positions, the top-left, top-right and bottom-left cells form one island. When (1, 1) arrives, two of its neighbours, (0, 1) and (1, 0), are land. Both are in the same island, because they connect through (0, 0). Only one merge should happen.

With the buggy class, the output is:

[1, 1, 1, 0]

The correct answer is [1, 1, 1, 1]. The count went to zero, which is impossible on a grid that has land in it. The if ra == rb: return line in _union is what stops this double subtraction. Remove it and the bug appears.

Building a grid of lists with repetition

The code in this lesson uses one flat list, so this mistake does not appear in it. It appears in many grid programs, so it is worth knowing. The expression [[False] * 3] * 2 looks like it makes two rows of three cells. It makes one row and puts two references to it in the outer list:

grid = [[False] * 3] * 2
grid[0][1] = True
print(grid)

[[False, True, False], [False, True, False]]

Changing row 0 also changed row 1, because both rows are the same list. A list comprehension makes a new row each time:

grid = [[False] * 3 for _ in range(2)]
grid[0][1] = True
print(grid)

[[False, True, False], [False, False, False]]

Practice

Exercise 1. Trace the algorithm by hand on a 2 by 2 grid with positions (0, 0), (1, 1) and (0, 1). Write the island count after each step.

Show the answer

After (0, 0) the count is 1. Cell (1, 1) has no land neighbours, so the count is 2. Cell (0, 1) has two land neighbours: (0, 0) to its left and (1, 1) below it. The new cell first adds 1, so the count is 3. The union with (0, 0) joins two islands, so the count is 2. The union with (1, 1) joins two islands, so the count is 1. The answer is [1, 2, 1]. The program prints this list.

Exercise 2. Explain in plain words why running a full flood fill after every addition is slower than union-find on large inputs.

Show the answer

A flood fill visits every land cell on the map, and it runs once per addition. If the grid has R cells and there are K additions, the flood fills do work proportional to R times K. Union-find only touches the new cell and its four neighbours on each addition, plus the short walks to the roots. That work does not depend on how much of the grid is already land.

Exercise 3. Change the code so that num_islands_2 also returns, after each addition, the size of the largest island. Describe the change and the lines you would edit.

Show the answer

Keep an attribute called self.largest, set to 0 in __init__. In add_land, after self.islands += 1 for the new cell, set self.largest to at least 1. In _union, after self.size[ra] += self.size[rb], set self.largest to the larger of its current value and self.size[ra]. Then return self.largest alongside the island count, for example as a tuple. Each union only needs to compare one number, so the work per addition stays the same.

Exercise 4. On a 1 by 5 grid, add the positions (0, 0), (0, 1), (0, 2), (0, 3), (0, 4) in that order. Write the island count after each step, and say which step causes the most merges.

Show the answer

The output is [1, 1, 1, 1, 1]. Each new cell joins the one island on its left. The count goes up by 1 for the new cell and down by 1 for the merge, so it stays at 1. Steps 2 to 5 each cause exactly one merge, so no step causes more than the others. Each new cell has only one land neighbour when it arrives, because the cell on its right is still water.

Exercise 5. A student writes this method, and it gives wrong counts when a position repeats. Find the bug and fix it.

def add_land(self, r, c):
    cell = self._id(r, c)
    self.land[cell] = True
    self.islands += 1
    # neighbour loop unchanged
    return self.islands

Show the answer

The method does not check whether the cell is already land. A repeated position adds 1 to self.islands again. Add the early return before anything else changes:

if self.land[cell]:
    return self.islands

Put it before self.land[cell] = True. With this check, a repeated position returns the current count without changing it.

Summary

• A new land cell adds 1 to the island count, then unions with each land neighbour.
• A union subtracts 1 only when it joins two different islands.
• Repeated positions must not change the count, so check for them first.
• The bounds check stops neighbours outside the grid from being read. Python does not raise an error for a negative index, so the check must cover negative numbers too.
• A slow flood fill on random small grids is a good way to check the fast method.