DataStructures.py

#################################### Union Find
class UnionFind:
    def __init__(self, nums):
        self.par = {}
        self.rank = {}
        for num in nums:
            self.par[num] = num
            self.rank[num] = 1
    
    def find(self, x):
        if x == self.par[x]:
            return x
        self.par[x] = self.find(self.par[x])
        return self.par[x]
    
    def union(self, x1, x2):
        p1, p2 = self.find(x1), self.find(x2)
        if p1 == p2:
            return False
        if self.rank[p1] > self.rank[p2]:
            self.par[p2] = p1
            self.rank[p1] += self.rank[p2]
        else:
            self.par[p1] = p2
            self.rank[p2] += self.rank[p1]
        return True
    
#################################### Segment Tree
class Node:
    def __init__(self, l, r, val=0, lnode=None, rnode=None):
        self.val = val
        self.left = lnode
        self.right = rnode
        self.L = l
        self.R = r

class SegmentTree:
    def __init__(self, nums):
        def _create(node):
            if node.L == node.R:
                node.val = nums[node.L]
            else:
                m = (node.L + node.R) // 2
                node.left = _create(Node(node.L, m))
                node.right = _create(Node(m+1, node.R))
                node.val = node.left.val + node.right.val
            return node
            
        self.root = _create(Node(0, len(nums)-1))

    def query(self, qleft, qright):
        def _query(head, q_l, q_r):
            if head.L > qright or head.R < qleft:
                return 0

            if (head.L >= qleft and head.R <= qright) or not head.left:
                return head.val
            
            return _query(head.left, q_l, q_r) + _query(head.right, q_l, q_r)
        
        return _query(self.root, qleft, qright)

    
    def update(self, index, value):
        def _updt(node, i, v):
            if node.L == node.R:
                node.val = v
                return v
            
            m = (node.L + node.R) // 2
            if m >= i:
                _updt(node.left, i, v)
            else:
                _updt(node.right, i, v)
                
            node.val = node.left.val + node.right.val
        
        _updt(self.root, index, value)
        
    def updateRange(self, qleft, qright, val):
        def _updt(head, ql, qr, val):
            if qr < head.L or ql > head.R:
                return

            if head.L >= ql and head.R <= qr:
                head.val += val
                if head.L != head.R: # Meaning its not a leaf node
                    _updt(head.left, ql, qr, val)
                    _updt(head.right, ql, qr, val)
                return

            _updt(head.left, ql, qr, val)
            _updt(head.right, ql, qr, val)
            head.val = head.left.val + head.right.val

        _updt(self.root, qleft, qright, val)
        
        
#################################### Binary Indexed Tree
class NumArray(object):
    def __init__(self, nums):
        self.n = len(nums)
        self.a, self.c = nums, [0] * (self.n + 1)
        for i in range(1, self.n+1):
            self.c[i] += nums[i-1]
            _ = i + (i & -i)
            if _ <= self.n:
                self.c[_] += self.c[i] 

    def update(self, i, val):
        d, self.a[i] = val - self.a[i], val
        i += 1
        while i <= self.n:
            self.c[i] += d
            i += (i & -i)

    def sumRange(self, i, j):
        return self.doSum(j+1, 0) - self.doSum(i, 0)
    
    def doSum(self, i, s):
        while i:
            s += self.c[i]
            i -= (i & -i)
        return s