During university I competed in quite a few programming contests including TopCoder, LeetCode, and ICPC. Over the years I collected quite a few recipes for functionalities almost built into python but needing just a litte work.

I’ve collected a few useful snippets here as a cookbook of some of my most frequently used recipes. Feel free to copy this code, but please consider:

  1. The rules of any contests/interviews you are participating in
  2. Unless you understand the whys of code, copying may actually be more work
Bulk Imports # 
import bisect
import heapq
import itertools
import math

from collections import defaultdict
from functools import lru_cache
from typing import *
from sortedcontainers import SortedDict, SortedList
ALPHABET # 
# ord('a') => 97
ALPHABET = 'abcdefghijklmnopqrstuvwxyz'

Searching #

Binary Search / Bisect / Partition # 
import bisect
import math

# First integer meeting predicate
def exposect_pred(pred, lo=1, limit=math.inf):
    hi = lo
    while hi < limit and not pred(hi):
        hi *= 2
    return bisect_pred(lo, hi, pred)

def bisect_pred(pred, lo, hi):
    while lo <= hi:
        mid = (lo + hi) // 2
        is_valid = pred(mid)
        
        if is_valid:
            hi = mid
            if lo == hi:
                return lo
        else:
            lo = mid + 1
    
    return -1 

# First index where value >= provided value
bisect.bisect_left([1, 1, 1, 2, 2, 2, 3, 3, 3], 2)
bisect.bisect_left([1, 1, 1, 2, 2, 2, 3, 3, 3], 1.9)
bisect.bisect_left([1, 1, 1, 2, 2, 2, 3, 3, 3], 1.1)
# returns 3                  ^

# First index where value > provided value
bisect.bisect_right([1, 1, 1, 2, 2, 2, 3, 3, 3], 2)
bisect.bisect_right([1, 1, 1, 2, 2, 2, 3, 3, 3], 2.1)
# returns 6                            ^

Data Structures #

Discrete Sets / Union Find / Union-Find / DS # 
class DisjointSets:
    def __init__(self):
        self.seen = set()
        self.parent = {}
        self.rank = {}
    def find(self, a):
        if a not in self.parent:
            return a
        ret = self.find(self.parent[a])
        self.parent[a] = ret
        return ret
    def union(self, a, b):
        self.seen.add(a)
        self.seen.add(b)
        a = self.find(a)
        b = self.find(b)
        if not a == b:
            rankA = self.rank.get(a, 0)
            rankB = self.rank.get(b, 0)
            if rankA > rankB:
                self.parent[b] = a
            elif rankB > rankA:
                self.parent[a] = b
            else:
                self.parent[b] = a
                self.rank[a] = rankA + 1
    def dsets(self):
        by_parent = defaultdict(set)
        for x in self.seen:
            by_parent[self.find(x)].add(x)
        return by_parent.values()

Graph #

Edge List to Undirected Map # 
from collections import defaultdict

# [ [ from, to, weight ] ] ==> { from: [(to, weight)], to: [(from, weight)] }
def edge_list_to_map(edges):
    edges_with = defaultdict(list)
    for from_id, to_id, weight in edges:
        # from_id, to_id = min(from_id, to_id), max(from_id, to_id)

        edges_with[from_id].append((to_id, weight))
        edges_with[to_id].append((from_id, weight))
    return edges_with
Dijkstra / BFS / Breadth First / A Star / A* # 
# Two modes:
# Provide t to receive cost to t
# Don't provide t to recieve min costs to all nodes
def dijkstra(edges_with, s, t=None):
    distance_to = {}
    q = [(0, s)]
    while q:
        dis, node = heapq.heappop(q)
        if node in distance_to:
            continue
        
        distance_to[node] = dis
        if t is not None and node == t:
            return dis
        
        for next_node, cost in edges_with[node]:
            assert cost > 0
            heapq.heappush(q, (dis + cost, next_node))
    
    return distance_to if t is None else None

Math #

PowerSet / Power Set / Set of Sets / All Sets # 
def powerset(iterable):
    s = list(iterable)
    return itertools.chain.from_iterable(
        itertools.combinations(s, r) for r in range(len(s)+1))

Func Tools #

Cached generator # 

from functools import lru_cache

def degenerate(fn):
    return lru_cache(lambda *args: list(fn(*args)))