Build a Union-FindStep 3 of 5

Union by size

12345

Step 3: Union by size

Always attaching one root under another can build tall, slow chains. Union by size keeps trees shallow: attach the smaller tree under the larger one.

Keep a size array where size[i] starts at 1. In union(a, b): find both roots; if equal, stop. Otherwise make the larger-size root the parent of the smaller-size root, and add the smaller size into the larger root's size.

Hints
Starter code
class DSU:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size = [1] * n

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

    def union(self, a, b):
        """Merge the sets containing a and b, attaching the smaller tree under the larger."""
        # TODO: find both roots; if equal, return. Otherwise make the larger-size root the
        # parent of the smaller-size root and add the smaller size into the larger.
        pass
Write it on a laptop

The editor and the test runner need a wide screen. Open this page on a desktop browser and your progress will be waiting.

PreviousNext step