M.T.Ramkrushna

Posted on Oct 3

14. Union-Find (DSU) + Minimum Spanning Trees

#ai #python #dsa

This is the natural next topic after graphs and shortest paths.

1. Union-Find / Disjoint Set Union (DSU)

Union-Find is used when you need to repeatedly answer:

“Are these two nodes in the same connected group?”

and merge groups together.

Core operations

find(x) → which group does x belong to?
union(a, b) → merge the groups containing a and b

2. Basic DSU

class DSU:
    def __init__(self, n):
        self.parent = list(range(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 = self.find(a)
        rb = self.find(b)

        if ra == rb:
            return False

        self.parent[rb] = ra
        return True

The important optimization here is:

self.parent[x] = self.find(self.parent[x])

This is path compression.

3. Union by Size / Rank

We can make DSU even faster by attaching the smaller tree to the larger tree.

class DSU:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size = [1] * 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 = self.find(a)
        rb = self.find(b)

        if ra == rb:
            return False

        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra

        self.parent[rb] = ra
        self.size[ra] += self.size[rb]

        return True

With path compression + union by size/rank:

Each operation is effectively O(α(n)), which is practically O(1).

4. When should you recognize DSU?

Think Union-Find when you see:

"connected components"
"merge groups"
"are these connected?"
"connect two nodes"
"redundant edge"
"number of provinces"
"accounts belong to same person"
"minimum spanning tree"

Mental trigger

Repeatedly connect things and check whether they're already connected → DSU.

5. Cycle Detection Using DSU

Suppose:

1 -- 2
2 -- 3
3 -- 1

When processing:

1-2
2-3
3-1

For 3-1:

find(3) == find(1)

They're already connected.

Therefore adding 3-1 creates a cycle.

Code

def has_cycle(n, edges):
    dsu = DSU(n)

    for u, v in edges:
        if not dsu.union(u, v):
            return True

    return False

Complexity:

O(E α(V))

Practically:

O(E)

6. Redundant Connection

Classic interview problem.

Given:

edges = [
    [1, 2],
    [2, 3],
    [3, 4],
    [1, 4],
    [1, 5]
]

The edge [1, 4] creates the cycle.

def findRedundantConnection(edges):
    n = len(edges)
    dsu = DSU(n + 1)

    for u, v in edges:
        if not dsu.union(u, v):
            return [u, v]

Pattern

For every edge:
    if endpoints already connected:
        cycle!
    else:
        union them

7. Number of Connected Components

Suppose:

0 -- 1     2 -- 3

4

There are 3 components.

We can count components using DSU.

def count_components(n, edges):
    dsu = DSU(n)
    components = n

    for u, v in edges:
        if dsu.union(u, v):
            components -= 1

    return components

Why?

Initially:

n nodes = n components

Every successful union reduces components by 1.

8. Minimum Spanning Tree (MST)

Now we move to an extremely important graph concept.

A spanning tree:

connects all vertices
contains no cycle
has exactly V - 1 edges

A Minimum Spanning Tree is the spanning tree with the minimum possible total edge weight.

Example:

     2
 A ------- B
 |         |
4|         |1
 |         |
 C ------- D
     3

We want to connect every node with minimum total cost.

9. MST vs Shortest Path

This distinction is extremely important.

Shortest Path

Question:

What's the cheapest way from A to B?

Algorithms:

BFS
Dijkstra
Bellman-Ford

Minimum Spanning Tree

Question:

What's the cheapest way to connect ALL nodes?

Algorithms:

Kruskal
Prim

Don't confuse them.

10. Kruskal's Algorithm

Kruskal is essentially:

Sort edges by weight and keep adding the cheapest edge that doesn't create a cycle.

DSU is perfect for detecting cycles.

Algorithm

1. Sort all edges by weight.
2. Create DSU.
3. For each edge from smallest to largest:
       if endpoints are in different components:
           add edge
           union endpoints
4. Stop after V - 1 edges.

11. Kruskal Python Template

def kruskal(n, edges):
    # edges = [(weight, u, v)]
    edges.sort()

    dsu = DSU(n)

    total = 0
    count = 0

    for weight, u, v in edges:
        if dsu.union(u, v):
            total += weight
            count += 1

            if count == n - 1:
                break

    return total

Example:

edges = [
    (1, 0, 1),
    (2, 1, 2),
    (3, 0, 2),
    (4, 2, 3)
]

print(kruskal(4, edges))

12. Kruskal Complexity

Sorting dominates:

O(E log E)

DSU operations are almost O(1).

Therefore:

Kruskal = O(E log E)

Good choice when the graph is represented naturally as an edge list, especially for sparse graphs.

13. Prim's Algorithm

Prim takes a different approach.

Instead of sorting all edges:

Start with one node and continuously add the cheapest edge that connects the current tree to a new node.

It looks similar to Dijkstra, but the goal is completely different.

Prim

Grow one MST.

Dijkstra

Find shortest paths from a source.

14. Prim Using a Min Heap

import heapq

def prim(n, graph):
    visited = [False] * n
    heap = [(0, 0)]   # (weight, node)

    total = 0
    count = 0

    while heap:
        weight, u = heapq.heappop(heap)

        if visited[u]:
            continue

        visited[u] = True
        total += weight
        count += 1

        for v, w in graph[u]:
            if not visited[v]:
                heapq.heappush(heap, (w, v))

    return total if count == n else -1

Graph:

graph = [
    [(1, 2), (2, 3)],
    [(0, 2), (2, 1)],
    [(0, 3), (1, 1)]
]

Each entry is:

(neighbor, weight)

15. Prim Complexity

Using adjacency lists + min heap:

O(E log V)

This is similar to Dijkstra's complexity.

16. Kruskal vs Prim

	Kruskal	Prim
Main idea	Cheapest edges globally	Grow one tree
Uses	DSU	Min heap
Input style	Edge list	Adjacency list
Complexity	O(E log E)	O(E log V)
Great for	Sparse graphs	Dense/adjacency-list graphs
Cycle detection	DSU	Visited set

Memorize

Kruskal → Sort + DSU
Prim    → Heap + Visited

That's one of the most useful interview shortcuts.

17. Important MST Properties

Property 1

An MST contains exactly:

V - 1 edges

Property 2

An MST cannot contain a cycle.

Property 3

For a connected graph, an MST connects every vertex.

Property 4

The MST does not necessarily give the shortest path between two particular vertices.

18. Cut Property

A useful theoretical concept:

For any cut in a graph, the minimum-weight edge crossing that cut can be part of an MST.

This is one of the ideas behind Prim and Kruskal.

For interviews, usually understanding the intuition is enough.

19. Very Important: MST ≠ Shortest Path Tree

Suppose:

A --1-- B
|       |
10      1
|       |
C --1-- D

MST tries to minimize:

TOTAL edge weight

Shortest path tries to minimize:

DISTANCE from a source

Different objective → potentially different edges.

20. Famous MST Problem: Min Cost to Connect All Points

Suppose you have points:

[(0,0), (2,2), (3,10), ...]

Cost to connect two points:

abs(x1-x2) + abs(y1-y2)

This is a classic MST problem.

You can solve it using either:

Prim

Kruskal

For Prim:

import heapq

def minCostConnectPoints(points):
    n = len(points)

    visited = [False] * n
    heap = [(0, 0)]

    total = 0
    used = 0

    while heap and used < n:
        cost, u = heapq.heappop(heap)

        if visited[u]:
            continue

        visited[u] = True
        used += 1
        total += cost

        x1, y1 = points[u]

        for v in range(n):
            if not visited[v]:
                x2, y2 = points[v]
                dist = abs(x1 - x2) + abs(y1 - y2)

                heapq.heappush(heap, (dist, v))

    return total

21. Interview Recognition Cheat Sheet

Problem wording	Think
Same connected component?	DSU
Merge groups	DSU
Detect redundant edge	DSU
Detect cycle in undirected graph	DSU
Number of components	DSU
Connect all nodes with minimum cost	MST
Minimum total connection cost	MST
Cheapest set of edges connecting everything	MST
Cheapest edge globally	Kruskal
Grow tree from a starting node	Prim
Shortest path from A to B	BFS/Dijkstra/etc.

22. The Big Graph Algorithm Map

At this point, your graph toolkit should look like this:

GRAPH
│
├── Traversal
│   ├── DFS
│   └── BFS
│
├── Connectivity
│   ├── DFS/BFS
│   └── DSU
│
├── Cycle Detection
│   ├── Undirected → DFS/DSU
│   └── Directed → DFS recursion stack
│
├── Shortest Path
│   ├── Unweighted → BFS
│   ├── 0/1 weights → 0-1 BFS
│   ├── Non-negative → Dijkstra
│   ├── Negative → Bellman-Ford
│   └── All pairs → Floyd-Warshall
│
├── Ordering
│   └── Topological Sort
│
└── Minimum Spanning Tree
    ├── Kruskal → Sort + DSU
    └── Prim → Heap

What to memorize

DSU:
find + union
path compression
union by size/rank

Kruskal:
sort edges
+ DSU

Prim:
min heap
+ visited

MST:
connect ALL vertices
with minimum TOTAL edge weight

Practice order

Redundant Connection → DSU
Number of Provinces → DSU
Number of Connected Components → DSU
Min Cost to Connect All Points → Prim/Kruskal
Kruskal MST implementation
Prim MST implementation

After this, the next major DSA topic is Dynamic Programming (DP) — one of the most important and commonly tested topics.