Build a Union-FindStep 4 of 5

Path compression and connected

12345

Step 4: Path compression and connected

Path compression flattens trees during lookups. Rewrite find(x): first walk to the root, then walk the path again pointing every visited node directly at the root. Future lookups become nearly constant time.

Also add connected(a, b): return True when a and b share a root, else False. This is the core query for connectivity problems.

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

    def find(self, x):
        """Return the root of x AND compress the path so visited nodes point at the root."""
        # TODO: walk to the root, then walk the path again setting each node to point directly
        # at the root. Return the root.
        pass

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

    def connected(self, a, b):
        """Return True if a and b are in the same set, else False."""
        # TODO: return whether find(a) equals find(b).
        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