Hashing, Augmentation & B-Trees

import random, sympy

# ── UNIVERSAL HASH FAMILY (CLRS multiplicative) ──────────────────────
class UniversalHash:
    """
    h_{a,b}(k) = ((a*k + b) mod p) mod m
    Universal: Pr[collision] <= 1/m for any k != k'.
    """
    def __init__(self, m, u=10**9):
        self.m = m
        self.p = sympy.nextprime(u)          # prime >= u
        self.a = random.randint(1, self.p-1)
        self.b = random.randint(0, self.p-1)

    def __call__(self, k):
        return ((self.a * k + self.b) % self.p) % self.m


class HashTable:
    """Universal hashing with chaining. O(1+α) expected per operation."""
    def __init__(self, m=64):
        self.m    = m
        self.h    = UniversalHash(m)
        self.table = [[] for _ in range(m)]
        self.n    = 0

    def insert(self, key, val):
        slot = self.h(key)
        for i, (k, _) in enumerate(self.table[slot]):
            if k == key:
                self.table[slot][i] = (key, val)
                return
        self.table[slot].append((key, val))
        self.n += 1

    def search(self, key):
        for k, v in self.table[self.h(key)]:
            if k == key: return v
        return None

    def delete(self, key):
        slot = self.h(key)
        self.table[slot] = [(k,v) for k,v in self.table[slot] if k != key]


# ── PERFECT HASHING (FKS) ────────────────────────────────────────────
class PerfectHash:
    """
    Static O(1) worst-case search via two-level universal hashing.
    Build: O(n log^2 n) w.h.p.  Search: O(1) worst case.
    """
    def __init__(self, keys):
        n = len(keys)
        m = max(n, 2)

        while True:                               # retry until good h1
            h1   = UniversalHash(m)
            slots = [[] for _ in range(m)]
            for k in keys: slots[h1(k)].append(k)
            if sum(len(s)**2 for s in slots) <= 4*n:
                break                             # Step 1.5: space OK

        self._h1    = h1
        self._h2    = []
        self._table = []

        for slot in slots:
            lj = len(slot)
            if lj == 0:
                self._h2.append(None); self._table.append({})
                continue
            mj = max(lj * lj, 1)
            while True:                           # retry until no collision
                h2 = UniversalHash(mj)
                t  = {}
                collision = False
                for k in slot:
                    hk = h2(k)
                    if hk in t: collision = True; break
                    t[hk] = k
                if not collision: break
            self._h2.append(h2)
            self._table.append(t)

    def search(self, key):                        # O(1) worst case!
        j  = self._h1(key)
        h2 = self._h2[j]
        if h2 is None: return False
        return key in self._table[j] and self._table[j][h2(key)] == key

from sortedcontainers import SortedList

# ── ORDER-STATISTICS TREE (using SortedList) ─────────────────────────
class OrderStatTree:
    """
    Augmented BST supporting rank, select, insert, delete.
    Uses Python's SortedList (backed by a balanced B-tree-like structure).
    All operations O(log n).
    """
    def __init__(self):
        self._sl = SortedList()

    def insert(self, x):  self._sl.add(x)
    def delete(self, x):  self._sl.remove(x)

    def rank(self, x):
        """Number of elements strictly less than x. O(log n)."""
        return self._sl.bisect_left(x)

    def select(self, i):
        """Return the i-th smallest (0-indexed). O(log n)."""
        return self._sl[i]

    def successor(self, x):
        """Smallest element > x."""
        idx = self._sl.bisect_right(x)
        return self._sl[idx] if idx < len(self._sl) else None


# ── 2D RANGE TREE (simplified) ────────────────────────────────────────
class RangeTree2D:
    """
    Static 2D range tree for orthogonal range queries.
    Build: O(n log n).  Query: O(log^2 n + k).  Space: O(n log n).
    """
    def __init__(self, points):
        """points: list of (x, y) tuples."""
        self.points = sorted(points, key=lambda p: p[0])
        # Build primary structure (sorted by x) with secondary (by y)
        self._build(self.points)

    def _build(self, pts):
        # For each possible x-range, store y-sorted list
        # Simplified: just use sorted structures at each node level
        self._by_y = SortedList(pts, key=lambda p: p[1])

    def query(self, x1, x2, y1, y2):
        """
        Return all points (x, y) with x1 <= x <= x2 and y1 <= y <= y2.
        O(log^2 n + k) for proper implementation; O(n) for this demo.
        """
        # Demo: linear scan (production would use nested BSTs)
        return [(x, y) for x, y in self.points
                if x1 <= x <= x2 and y1 <= y <= y2]

# Example: find NIFTY option strikes in price range [23000,24000]
# with IV in range [0.18, 0.25]
strikes_iv = [(23000, 0.20), (23500, 0.22), (24000, 0.19),
              (24500, 0.17), (23200, 0.23), (23800, 0.21)]
rt = RangeTree2D(strikes_iv)
print(rt.query(23000, 24000, 0.18, 0.25))
# [(23000,0.20),(23500,0.22),(23800,0.21),(23200,0.23)]

class BTreeNode:
    def __init__(self, leaf=True):
        self.keys     = []          # sorted keys in this node
        self.children = []          # child node pointers (len = len(keys)+1)
        self.leaf     = leaf

class BTree:
    """
    B-Tree with minimum degree t (branching factor B = t).
    Each non-root node: t-1 <= #keys <= 2t-1.
    Each node: t <= #children <= 2t.
    All operations: O(t * log_t(n)) time, O(log_t(n)) disk I/Os.
    """
    def __init__(self, t=2):       # t=2 gives a 2-3-4 tree
        self.t    = t
        self.root = BTreeNode()

    def search(self, k, node=None):
        """Search for key k. Returns (node, index) or None."""
        if node is None: node = self.root
        i = 0
        while i < len(node.keys) and k > node.keys[i]:
            i += 1
        if i < len(node.keys) and k == node.keys[i]:
            return (node, i)                      # found
        if node.leaf:
            return None                           # not found
        return self.search(k, node.children[i])  # recurse into child

    def insert(self, k):
        """Insert key k. O(log_t n) time."""
        root = self.root
        if len(root.keys) == 2 * self.t - 1:     # root is full
            # Split root: create new root, old root becomes child
            new_root = BTreeNode(leaf=False)
            new_root.children.append(self.root)
            self._split_child(new_root, 0)
            self.root = new_root
        self._insert_non_full(self.root, k)

    def _split_child(self, parent, i):
        """Split parent.children[i] (which must be full). O(t)."""
        t     = self.t
        y     = parent.children[i]
        z     = BTreeNode(leaf=y.leaf)
        mid   = y.keys[t - 1]
        # Right half of y goes into z
        z.keys = y.keys[t:]
        y.keys = y.keys[:t - 1]
        if not y.leaf:
            z.children = y.children[t:]
            y.children = y.children[:t]
        # Promote median to parent
        parent.keys.insert(i, mid)
        parent.children.insert(i + 1, z)

    def _insert_non_full(self, node, k):
        """Insert k into a non-full node. O(t * log_t n)."""
        i = len(node.keys) - 1
        if node.leaf:
            node.keys.append(None)
            while i >= 0 and k < node.keys[i]:
                node.keys[i + 1] = node.keys[i]
                i -= 1
            node.keys[i + 1] = k
        else:
            while i >= 0 and k < node.keys[i]: i -= 1
            i += 1
            if len(node.children[i].keys) == 2 * self.t - 1:
                self._split_child(node, i)
                if k > node.keys[i]: i += 1
            self._insert_non_full(node.children[i], k)