troupe-standup/wizard-work/beam/robin_hood_hash.py

"""Robin Hood open-addressing hash table.

Keys and values can be any hashable/anything. Tombstones are NOT used; deletion
re-hashes the run tail to keep probe distances valid (backward-shift deletion).

Invariant: for every occupied slot s with key k inserted at home h = hash(k) % n,
probe_distance(s) = (s - h) mod n, and the table is "Robin Hood": slots are
ordered so that no slot has a strictly greater probe distance than an earlier
slot in the same probe chain (Robin Hood property is maintained on insert by
swapping the new entry with the incumbent when the new entry's probe distance
exceeds the incumbent's).

Complexities: amortized O(1) insert/lookup/delete at load factor < ~0.75.
Resizes by doubling when load factor exceeds 0.7.
"""
from __future__ import annotations


class RobinHoodHash:
    def __init__(self, capacity: int = 16):
        if capacity < 4:
            capacity = 4
        # round up to power of two for fast modulo
        c = 1
        while c < capacity:
            c <<= 1
        self._cap = c
        self._keys = [None] * self._cap
        self._vals = [None] * self._cap
        self._dists = [0] * self._cap  # 0 means empty (probe dist >=1 when occupied)
        self._size = 0

    def _home(self, key):
        return hash(key) & (self._cap - 1)

    def _resize(self, new_cap):
        old_keys = self._keys
        old_vals = self._vals
        self._cap = new_cap
        self._keys = [None] * new_cap
        self._vals = [None] * new_cap
        self._dists = [0] * new_cap
        self._size = 0
        for k, v in zip(old_keys, old_vals):
            if k is not None or (k is None and old_keys.index(k) is not None):
                pass
        # simpler: track occupied via dists of old arrays
        # Rebuild from old arrays captured before overwrite
        pass

    def __len__(self):
        return self._size

    def _insert_raw(self, key, value):
        idx = self._home(key)
        dist = 1
        while True:
            if self._dists[idx] == 0:
                self._keys[idx] = key
                self._vals[idx] = value
                self._dists[idx] = dist
                self._size += 1
                return
            # Robin Hood: if current entry's probe distance < dist, swap
            if self._dists[idx] < dist:
                # swap
                self._keys[idx], key = key, self._keys[idx]
                self._vals[idx], value = value, self._vals[idx]
                self._dists[idx], dist = dist, self._dists[idx]
                # now continue inserting the displaced entry
            idx = (idx + 1) & (self._cap - 1)
            dist += 1

    def __setitem__(self, key, value):
        # update existing
        idx = self._home(key)
        dist = 1
        while self._dists[idx] != 0 and dist <= self._cap:
            if dist > self._dists[idx]:
                break
            if self._keys[idx] == key:
                self._vals[idx] = value
                return
            idx = (idx + 1) & (self._cap - 1)
            dist += 1
        # not found -> insert (may resize)
        if (self._size + 1) / self._cap > 0.7:
            self._grow()
        self._insert_raw(key, value)

    def __getitem__(self, key):
        idx = self._home(key)
        dist = 1
        while self._dists[idx] != 0:
            if dist > self._dists[idx]:
                raise KeyError(key)
            if self._keys[idx] == key:
                return self._vals[idx]
            idx = (idx + 1) & (self._cap - 1)
            dist += 1
            if dist > self._cap:
                raise KeyError(key)
        raise KeyError(key)

    def __contains__(self, key):
        try:
            self.__getitem__(key)
            return True
        except KeyError:
            return False

    def __delitem__(self, key):
        # find slot
        idx = self._home(key)
        dist = 1
        found = None
        while self._dists[idx] != 0:
            if dist > self._dists[idx]:
                raise KeyError(key)
            if self._keys[idx] == key:
                found = idx
                break
            idx = (idx + 1) & (self._cap - 1)
            dist += 1
            if dist > self._cap:
                break
        if found is None:
            raise KeyError(key)
        # backward-shift deletion
        self._keys[found] = None
        self._vals[found] = None
        self._dists[found] = 0
        self._size -= 1
        # rehash following run
        j = (found + 1) & (self._cap - 1)
        while self._dists[j] != 0:
            k = self._keys[j]
            v = self._vals[j]
            self._keys[j] = None
            self._vals[j] = None
            self._dists[j] = 0
            self._size -= 1
            self._insert_raw(k, v)
            j = (j + 1) & (self._cap - 1)

    def _grow(self):
        old_keys = list(self._keys)
        old_vals = list(self._vals)
        old_dists = list(self._dists)
        self._cap *= 2
        self._keys = [None] * self._cap
        self._vals = [None] * self._cap
        self._dists = [0] * self._cap
        self._size = 0
        for i in range(len(old_keys)):
            if old_dists[i] != 0:
                self._insert_raw(old_keys[i], old_vals[i])

    def keys(self):
        return [k for k, d in zip(self._keys, self._dists) if d != 0]

    def items(self):
        return [(k, v) for k, v, d in zip(self._keys, self._vals, self._dists) if d != 0]