Data Structures

min-heap.py

class MinHeap:
    """A min-heap using an array"""
    # storing tree in level-order
    arr = []
    @staticmethod
    def parent(i): return (i - 1) // 2
    @staticmethod
    def children(i): return (2*i + 1, 2*i + 2)

    def insert(self, a):
        self.arr.append(a)
        i = len(self.arr) - 1
        p = self.parent(i)
        while p >= 0 and self.arr[i] < self.arr[p]:
            self.arr[i], self.arr[p] = self.arr[p], self.arr[i]
            i = p
            p = self.parent(i)

    def delete(self, a):
        """assuming arr has distinct elements"""
        i, l = heap.arr.index(a), len(self.arr)
        self.arr[i], self.arr[l-1] = self.arr[l-1], self.arr[i]
        self.arr.pop()
        self._heapify(i)

    def peek(self):
        if len(self.arr) > 0:
            return self.arr[0]
        
    def extract_min(self):
        min_ self.arr[0]
        last = len(self.arr) - 1
        self.arr[0] = self.arr[last]
        del self.arr[last]
        self._heapify(0)
        return min_

    def _heapify(self, i):
        left, right = self.children(i)
        smallest = i

        if left < len(self.arr) and self.arr[left] < self.arr[smallest]:
            smallest = left
        if right < len(self.arr) and self.arr[right] < self.arr[smallest]:
            smallest = right

        if smallest != i:
            self.arr[i], self.arr[smallest] = self.arr[smallest], self.arr[i]
            self._heapify(smallest)


if __name__ == '__main__':
    heap = MinHeap()
    heap.insert(7)
    heap.insert(2)
    heap.insert(5)
    heap.insert(1)
    heap.insert(9)
    heap.insert(3)
    print(heap.extract_min())  # 1
    print(heap.peek())  # 3
    heap.delete(3)
    heap.delete(9)
    print(heap.arr)

prefix-sum-1d.py

if __name__ == '__main__':
    nums = [7, 2, 5, 10, 8]
    prefix = [0] + nums
    for i in range(1, len(prefix)):
        prefix[i] += prefix[i-1]
        
    print(prefix)  # [0, 7, 9, 14, 24, 32]
    print(nums[1:4])  # [2, 5, 10]
    print(prefix[4] - prefix[1])  # 17

python-priority-queue.py

"""
Priority Queue with modification and deletion semantics.

https://docs.python.org/3/library/heapq.html#priority-queue-implementation-notes
"""
from heapq import heappush, heappop

pq = []
entry_dict = {}  # O(1) access of any task's entry
count = 0


def add(task, priority):
    global count
    # count acts as tie-breaker between two tasks with the same priority
    # a global counter ensures tasks are pop-ed with sort stability
    count += 1
    entry = [priority, count, task]
    entry_dict[task] = entry
    heappush(pq, entry)


def remove(task):
    entry = entry_dict.pop(task)
    entry[-1] = 'REMOVED'


def update(task, priority):
    if task in entry_dict:
        remove(task)
    add(task, priority)


def pop():
    while pq:
        _, _, task = heappop(pq)
        if task != 'REMOVED':
            del entry_dict[task]
            return task


if __name__ == '__main__':
    add('t1', 1)
    add('t2', 2)
    add('t3', 2)
    add('t4', 3)

    print(pop())  # t1
    print(pop())  # t2 (Sort stability: t2 and t3 has same priority but t2 was added first)

    add('t5', 4)
    update('t4', 1)

    print(pop())  # t4 (since its priority was updated to 1)

trie.py

"""Trie for storing words"""

class Node:
    """Represents a letter"""
    def __init__(self):
        # letter -> Node
        self.children = {}
        # is it the last letter of a word
        self.is_last = False

    def attach(self, word):
        cur = self
        for i, letter in enumerate(word):
            node = cur.children.get(letter, Node())
            cur.children[letter] = node
            node.is_last = i == len(word)-1
            cur = node

            
if __name__ == '__main__':
    root = Node()
    root.attach('blue')
    root.attach('bleak')
    root.attach('sit')
    root.attach('sea')

union-find.py

"""
https://www.hackerearth.com/practice/notes/disjoint-set-union-union-find/
Above link also includes:
- weighted union
- union with path compression
"""
from collections import defaultdict


class UnionFind:
    def __init__(self, items):
        self.parent = {item: None for item in items}
    
    def _root(self, x):
        if self.parent[x] is None:
            return x
        return self._root(self.parent[x])
    
    def union(self, x, y):
        root_x = self._root(x)
        root_y = self._root(y)
        if root_x != root_y:  # if x and y are already not connected
            self.parent[root_x] = root_y 
        
    def find(self, x):
        return self._root(x)
    
    def groups(self):  # connected components
        grps = defaultdict(list)
        for item in self.parent.keys():
            grps[self._root(item)].append(item)
        return grps


if __name__ == '__main__':
    uf = UnionFind(list(range(1, 7)))
    uf.union(5, 1)
    uf.union(2, 4)
    uf.union(3, 5)
    uf.union(2, 6)
    uf.union(6, 2)
    # the above edges will form the forest:
    #  1 - 3 - 5
    #  2 - 4 - 6

    # all have same parent, hence connected:
    print(uf.find(1), uf.find(3), uf.find(5))  # 1 1 1
    print(uf.find(2), uf.find(4), uf.find(6))  # 6 6 6
    print(uf.groups())  # defaultdict(<class 'list'>, {1: [5, 1, 3], 6: [2, 4, 6]})