Union-Find: number of provinces

You have n cities and an n x n matrix where is_connected[i][j] == 1 means cities i and j are directly linked. A province is a group of cities reachable from one another (directly or through intermediaries). Count the provinces.

This is a "how many connected groups?" question — the home turf of Union-Find (Disjoint Set Union). Every city starts as its own group. Each time two cities are connected, you union their groups. At the end, the number of distinct roots is the number of provinces.

parent = [0, 1, 2]   # each city is its own root to begin with
# connect 0 and 1 -> they now share a root -> 3 groups become 2

Two tricks keep it fast:

• find with path compression — when you walk up to a root, re-point each node directly at the root so future lookups are flat.
• union — point one root at the other; if they already share a root, do nothing.

def find(x):
    while parent[x] != x:
        parent[x] = parent[parent[x]]   # path compression
        x = parent[x]
    return x

Your task

Write find_circle_num(is_connected) returning the number of provinces. is_connected is a square matrix; is_connected[i][i] is always 1, and the matrix is symmetric.

find_circle_num([[1,1,0],[1,1,0],[0,0,1]])   # 2   (0-1 together, 2 alone)
find_circle_num([[1,0,0],[0,1,0],[0,0,1]])   # 3   (nobody connected)
find_circle_num([[1,1,1],[1,1,1],[1,1,1]])   # 1   (all one province)
find_circle_num([[1]])                       # 1

Hint: build a parent array, union every i, j where is_connected[i][j] is 1, then count how many cities are their own root.