Implements various algorithms for union find, based on Algorithms, Part I

union-find.py

class QuickFind(object):
    """
    Implements the quick find algorithm for the dynamic connectivity problem
    
    Cost model (in array read/writes):
        Initialization: N
        Union: N
        Find: 1
    
    Algorithm:
        Start with an array N elements long (corresponding the number of objects)
        Initialize the array with 0...N-1
        To connect (union) two objects, set their ids to be the same.
        All connected objects will therefore have the same id. 
        
    >>> qf = QuickFind(5)
    >>> qf.union(1, 2)
    >>> qf.union(3, 4)
    >>> qf.union(3, 2)
    >>> qf.find(1,4)
    True
    >>> qf.find(1, 0)
    False
    """
    def __init__(self, n):
        self.id = [x for x in range(n)]
    def union(self, p, q):
        pid = self.id[p]
        qid = self.id[q]
        for x in range(len(self.id)):
            if self.id[x] == pid:
                self.id[x] = qid
    def find(self, p, q):
        return self.id[p] == self.id[q]
        
        
class QuickUnion(object):
    """
    Implements the quick union algorithm for the dynamic connectivity problem
    
    Cost model (in array read/writes):
        Initialization: N
        Union: N (depends on depth to root)
        Find: N
    
    Algorithm:
        Start with an array N elements long (corresponding the number of objects)
        Initialize the array with 0...N-1
        id[i] is the root of i
        To connect (union) two objects, set their root ids to be the same.
        All connected objects will therefore have the same root id.   
        
    >>> qu = QuickUnion(5)
    >>> qu.union(1,2)
    >>> qu.union(3,4)
    >>> qu.union(3,2)
    >>> qu.find(1,4)
    True
    >>> qu.find(1,0)
    False
    """
    def __init__(self, n):
        self.id = [x for x in range(n)]
    def union(self, p, q):
        proot = self._root(p)
        qroot = self._root(q)
        self.id[proot] = qroot
    def find(self, p, q):
        return self._root(p) == self._root(q)
    def _root(self, i):
        while self.id[i] is not i:
            i = self.id[i]
        return i
        
class WeightedQuickUnion(object):
    """
    Implements the quick union algorithm for the dynamic connectivity problem.
    Considers the 'weight' (size) of each tree before merging.
    
    Cost model (in array read/writes):
        Initialization: N
        Union: lg N
        Find: lg N
    
    Algorithm:
        Start with an array N elements long (corresponding the number of objects)
        Initialize the array with 0...N-1
        id[i] is the root of i
        To connect (union) two objects, set their root ids to be the same.
        Always merge the smaller tree under the larger tree (faster because _root is called less)
        All connected objects will therefore have the same root id.   
        
    >>> wqu = WeightedQuickUnion(5)
    >>> wqu.union(1,2)
    >>> wqu.union(3,4)
    >>> wqu.union(3,2)
    >>> wqu.find(1,4)
    True
    >>> wqu.find(1,0)
    False
    """
    def __init__(self, n):
        self.id = [x for x in range(n)]
        self.sz = [1] * n
    def union(self, p, q):
        proot = self._root(p)
        qroot = self._root(q)
        if proot == qroot:
            return
        if self.sz[proot] < self.sz[qroot]:
            self.id[proot] = qroot
            self.sz[qroot] += self.sz[proot]
        else:
            self.id[qroot] = proot
            self.sz[proot] += self.sz[qroot]
    def find(self, p, q):
        return self._root(p) == self._root(q)
    def _root(self, i):
        while self.id[i] is not i:
            i = self.id[i]
        return i
        
class WeightedQuickUnionPC(object):
    """
    Implements the quick union algorithm for the dynamic connectivity problem.
    Considers the 'weight' (size) of each tree before merging.
    Has path compression.
    
    Cost model (in array read/writes):
        Initialization: N
        Union: lg* N (essentially linear)
        Find: lg* N (essentially linear)
    
    Algorithm:
        Start with an array N elements long (corresponding the number of objects)
        Initialize the array with 0...N-1
        id[i] is the root of i
        To connect (union) two objects, set their root ids to be the same.
        Always merge the smaller tree under the larger tree (faster because _root is called less)
        Compress paths as we find roots - set the root of every other parent to its parent
        All connected objects will therefore have the same root id.   
        
    >>> wqupc = WeightedQuickUnionPC(5)
    >>> wqupc.union(1,2)
    >>> wqupc.union(3,4)
    >>> wqupc.union(3,2)
    >>> wqupc.find(1,4)
    True
    >>> wqupc.find(1,0)
    False
    """
    def __init__(self, n):
        self.id = [x for x in range(n)]
        self.sz = [1] * n
    def union(self, p, q):
        proot = self._root(p)
        qroot = self._root(q)
        if proot == qroot:
            return
        if self.sz[proot] < self.sz[qroot]:
            self.id[proot] = qroot
            self.sz[qroot] += self.sz[proot]
        else:
            self.id[qroot] = proot
            self.sz[proot] += self.sz[qroot]
    def find(self, p, q):
        return self._root(p) == self._root(q)
    def _root(self, i):
        while self.id[i] is not i:
            self.id[i] = self.id[self.id[i]] # path compression
            i = self.id[i]
        return i