DSA Number of Provinces

The problem

Number of Provinces gives a square matrix about n cities. Entry i, j is one when cities i and j link directly, and zero when they do not.

A province is a maximal group of cities linked directly or through others. Maximal means no linked city sits outside the group. The answer is the number of provinces.

cases = [
    ([[1, 1, 0], [1, 1, 0], [0, 0, 1]], 2),
    ([[1, 0, 0, 1], [0, 1, 1, 0], [0, 1, 1, 0], [1, 0, 0, 1]], 2),
    ([[1, 0, 0], [0, 1, 0], [0, 0, 1]], 3),
]
for matrix, expected in cases:
    print("map", matrix, "-> provinces:", expected)

The second map holds two pairs, cities 0 and 3, and cities 1 and 2. No link joins the pairs. The third map links nothing, so every city stands alone.

Intuition: subtract the links from the cities

The first idea counts links. Each link seems to fuse two cities into one group. Cities minus links looks like a province count.

matrix = [[1, 1, 1], [1, 1, 1], [1, 1, 1]]
n = 3
links = sum(matrix[i][j] for i in range(n) for j in range(i + 1, n))
print("cities:", n, "links:", links)
print("cities minus links:", n - links)
print("true provinces: 1")

The count reaches zero provinces for a map that plainly holds one. The third link joins two cities already in the same group. Only merges that join different groups lower the count.

Brute force: flood each province

The brute force floods outward from the first unvisited city. A flood is a walk that spreads along links to every reachable city. Every city the flood reaches belongs to one province, and the next unvisited city starts the next flood.

A flood scans the full row of each city it visits. The check counter records every matrix entry read.

def provinces_brute(matrix):
    n = len(matrix)
    seen = [False] * n
    checks = 0
    count = 0
    for start in range(n):
        if seen[start]:
            continue
        count += 1
        stack = [start]
        seen[start] = True
        while stack:
            city = stack.pop()
            for other in range(n):
                checks += 1
                if matrix[city][other] == 1 and not seen[other]:
                    seen[other] = True
                    stack.append(other)
    return count, checks

for matrix in ([[1, 1, 0], [1, 1, 0], [0, 0, 1]], [[1, 0, 0], [0, 1, 0], [0, 0, 1]]):
    count, checks = provinces_brute(matrix)
    print("provinces:", count, "matrix checks:", checks)

Each flood reads n entries per city, and the floods together read the whole matrix. The method is simple and correct, and it pays the full matrix price.

The optimal solution: union-find over the links

Union-find keeps one set per province, and each link unions the sets of its two cities while the province count falls only when a union really merges. That one sentence is the whole method.

Find returns the root that names a set. Two cities share a province exactly when their roots match. The final answer is the number of distinct roots.

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n

    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        if self.rank[ra] < self.rank[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra
        if self.rank[ra] == self.rank[rb]:
            self.rank[ra] += 1
        return True

def provinces(matrix):
    n = len(matrix)
    uf = UnionFind(n)
    merges = 0
    for i in range(n):
        for j in range(i + 1, n):
            if matrix[i][j] == 1:
                if uf.union(i, j):
                    merges += 1
    roots = len({uf.find(i) for i in range(n)})
    return roots, merges

cases = [
    ([[1, 1, 0], [1, 1, 0], [0, 0, 1]], 2),
    ([[1, 0, 0, 1], [0, 1, 1, 0], [0, 1, 1, 0], [1, 0, 0, 1]], 2),
    ([[1, 0, 0], [0, 1, 0], [0, 0, 1]], 3),
]
for matrix, expected in cases:
    count, merges = provinces(matrix)
    print("map -> provinces:", count, "merges:", merges, "match:", count == expected)

All three maps pass. Map two needs two merges for its two pairs, and the counter shows them.

Walkthrough: two pairs in the second map

Four cities start as four single sets. The matrix holds two links. Each link triggers one union, and the parent array prints after each one.

parent = list(range(4))

def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])
    return parent[x]

def union(a, b):
    ra, rb = find(a), find(b)
    if ra != rb:
        parent[rb] = ra

print("start:", parent)
union(0, 3)
print("after union(0, 3):", parent)
union(1, 2)
print("after union(1, 2):", parent)
print("roots:", sorted({find(i) for i in range(4)}))
print("provinces:", len({find(i) for i in range(4)}))

Two roots remain standing. Each root names one province, so the answer is two. City 3 shares a root with city 0 without any direct scan of their link beyond the one union.

Complexity

The flood reads every matrix entry, so its time is O(n squared). Union-find reads only the upper triangle, half the entries, and each find is nearly constant with path compression.

Both methods store one small array per city, so their extra space is O(n).

def chain(n):
    m = [[0] * n for _ in range(n)]
    for i in range(n):
        m[i][i] = 1
        if i + 1 < n:
            m[i][i + 1] = m[i + 1][i] = 1
    return m

def brute_checks(matrix):
    n = len(matrix)
    seen = [False] * n
    checks = 0
    for start in range(n):
        if seen[start]:
            continue
        stack = [start]
        seen[start] = True
        while stack:
            city = stack.pop()
            for other in range(n):
                checks += 1
                if matrix[city][other] == 1 and not seen[other]:
                    seen[other] = True
                    stack.append(other)
    return checks

def union_checks(matrix):
    n = len(matrix)
    parent = list(range(n))

    def find(x):
        if parent[x] != x:
            parent[x] = find(parent[x])
        return parent[x]

    checks = 0
    for i in range(n):
        for j in range(i + 1, n):
            checks += 1
            if matrix[i][j] == 1:
                ra, rb = find(i), find(j)
                if ra != rb:
                    parent[rb] = ra
    return checks

for n in (4, 8, 16):
    matrix = chain(n)
    print("cities", n, "- flood checks:", brute_checks(matrix), " union checks:", union_checks(matrix))

At 16 cities the flood pays 256 checks against 120 for the triangle scan. The matrix itself forces real work on any method. Reading only half of it is the saving on offer.

Edge cases and pitfalls

One city forms one province. A lone city needs no links to count.

The diagonal of the matrix is always one. Reading it as links would union every city with itself. Those unions change nothing, and the scan skips the diagonal for that reason.

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n

    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        if self.rank[ra] < self.rank[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra
        if self.rank[ra] == self.rank[rb]:
            self.rank[ra] += 1
        return True

def provinces(matrix):
    n = len(matrix)
    uf = UnionFind(n)
    merges = 0
    for i in range(n):
        for j in range(i + 1, n):
            if matrix[i][j] == 1:
                if uf.union(i, j):
                    merges += 1
    roots = len({uf.find(i) for i in range(n)})
    return roots, merges

print("one city:", provinces([[1]])[0])
print("two unlinked cities:", provinces([[1, 0], [0, 1]])[0])
print("all linked:", provinces([[1, 1, 1], [1, 1, 1], [1, 1, 1]])[0])
print("empty matrix:", provinces([])[0])

The empty matrix answers zero. No cities means no provinces, and the loops simply never run.