Disjoint Set Implementation

disjoint_set_implementation.py

class UnionFind:
    def __init__(self):
        self.parent = {}
        self.rank = {}

    def add(self, x):
        """
        Adds a new element to the Union-Find structure.
        Initializes the element as its own parent (singleton set).
        """
        if x not in self.parent:
            self.parent[x] = x  # Parent of itself
            self.rank[x] = 0    # Rank is initially 0

    def find(self, x):
        """
        Finds the representative (root) of the set containing x.
        Uses path compression for optimization.
        """
        if x not in self.parent:
            raise ValueError(f"Element {x} is not in the structure. Use 'add' to add it first.")
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])  # Path compression
        return self.parent[x]

    def union(self, x, y):
        """
        Merges the sets containing x and y.
        Uses union by rank for optimization.
        """
        root_x = self.find(x)
        root_y = self.find(y)

        if root_x != root_y:
            # Union by rank
            if self.rank[root_x] > self.rank[root_y]:
                self.parent[root_y] = root_x
            elif self.rank[root_x] < self.rank[root_y]:
                self.parent[root_x] = root_y
            else:
                self.parent[root_y] = root_x
                self.rank[root_x] += 1

    def connected(self, x, y):
        """
        Checks if two elements are in the same set.
        """
        return self.find(x) == self.find(y)