# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Hash table load factor: measure the probes.

A linear-probing table of 2^20 slots filled to load factors from 0.5 to
0.95; average probes for successful and failed lookups against the formulas
1/2 (1 + 1/(1 - a)) and 1/2 (1 + 1/(1 - a)^2). Then lookup time for linear
probing against chaining at the same loads.
"""
import argparse, json, random, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--size", type=int, default=1 << 20, help="Slots in the table.")
p.add_argument("--alphas", type=lambda s: [float(v) for v in s.split(",")], default=[0.5, 0.7, 0.8, 0.9, 0.95], help="Load factors.")
p.add_argument("--lookups", type=int, default=100_000, help="Lookups per measurement.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
random.seed(args.seed)
size = args.size
keys = [random.getrandbits(64) for _ in range(int(size * 0.96))]
table, n = [None] * size, 0
chains = [[] for _ in range(size)]
rows = []
for alpha in args.alphas:
    while n < alpha * size:
        k = keys[n]; i = k % size
        while table[i] is not None: i = (i + 1) % size
        table[i] = k; chains[k % size].append(k); n += 1
    def probes(key):
        i, c = key % size, 1
        while table[i] is not None and table[i] != key: i = (i + 1) % size; c += 1
        return c
    present = [keys[random.randrange(n)] for _ in range(args.lookups)]       # index, do not copy the list
    absent = [random.getrandbits(64) for _ in range(args.lookups)]
    hit = sum(map(probes, present)) / args.lookups; miss = sum(map(probes, absent)) / args.lookups
    t = time.perf_counter(); [probes(k) for k in present]; t_lin = (time.perf_counter() - t) / args.lookups * 1e9
    t = time.perf_counter(); [k in chains[k % size] for k in present]; t_ch = (time.perf_counter() - t) / args.lookups * 1e9
    rows.append({"alpha": alpha, "hit_probes": hit, "miss_probes": miss, "hit_formula": 0.5 * (1 + 1 / (1 - alpha)),
                 "miss_formula": 0.5 * (1 + 1 / (1 - alpha) ** 2), "inv": 1 / (1 - alpha), "linear_ns": t_lin, "chain_ns": t_ch})
    r = rows[-1]
    print(f"alpha={alpha:.2f}  hit {hit:6.2f} (formula {r['hit_formula']:6.2f})  miss {miss:7.2f} (formula {r['miss_formula']:7.2f})   "
          f"lookup: linear {t_lin:5.0f} ns, chaining {t_ch:5.0f} ns")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"loads": rows}}, f)