Z-order floor

# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""The Z-order floor.

A million random 2-D points split into U = 4,096 units under three layouts:
sorted on dimension 1, Z-order (bits of the two coordinates interleaved) and
random. Counts the units whose min/max box overlaps a thin slab query on
dimension 2 alone, and on dimension 1 alone; then the cost of a 50/50 mix,
against √U and 0.5·U.
"""
import argparse, json
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--units", type=int, default=4096, help="Units U.")
p.add_argument("--points", type=int, default=1 << 20, help="Points.")
p.add_argument("--queries", type=int, default=50, help="Queries per dimension.")
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()

rng = np.random.default_rng(args.seed); pts = rng.integers(0, 1 << 16, (args.points, 2)).astype(np.uint64)
def z(x):
    out = np.zeros(len(x), dtype=np.uint64)
    for b in range(16):
        out |= ((x[:, 0] >> np.uint64(b)) & np.uint64(1)) << np.uint64(2 * b); out |= ((x[:, 1] >> np.uint64(b)) & np.uint64(1)) << np.uint64(2 * b + 1)
    return out
orders = {"sorted on dimension 1": np.lexsort((pts[:, 1], pts[:, 0])), "Z-order": np.argsort(z(pts), kind="stable"), "random": rng.permutation(args.points)}
rows = []
for name, o in orders.items():
    u = pts[o].reshape(args.units, -1, 2); mn, mx = u.min(axis=1), u.max(axis=1)
    cost = {}
    for dim in (0, 1):
        vals = rng.integers(0, 1 << 16, args.queries)
        cost[dim] = float(np.mean([((mn[:, dim] <= v) & (v <= mx[:, dim])).sum() for v in vals]))      # a very thin slab: one value
    rows.append({"layout": name, "dim1_query": cost[0], "dim2_query": cost[1], "mix_50_50": (cost[0] + cost[1]) / 2})
summary = {"sqrt_u": args.units ** 0.5, "half_u": args.units / 2}
for r in rows: print(f"{r['layout']:<22} dimension-1 query {r['dim1_query']:7.1f} units   dimension-2 query {r['dim2_query']:7.1f}   50/50 mix {r['mix_50_50']:7.1f}")
print(f"√U = {summary['sqrt_u']:.0f}, 0.5·U = {summary['half_u']:.0f}")

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