Quick and easy way convert an array index into an index into a permutation of the array

shuffled_index.py

import numpy as np


def num_bits(max_i: int) -> int:
    return len(bin(max_i)[2:])


def int_to_bits(i: int, n: int, dtype=int) -> list[int]:
    return [dtype(int(b)) for b in bin(i)[2:].zfill(n)]


def bits_to_int(bits: list[Union[int, bool]]) -> int:
    out = 0
    for bit in bits:
        out = (out << 1) | int(bit)
    return out


def shuffled_index(i: int, max_i: int, seed: int) -> int:
    assert i <= max_i
    n = num_bits(max_i)

    def init_transform(s: int, size: tuple[int, ...]) -> np.ndarray:
        rng = np.random.default_rng(s)
        return rng.choice([False, True], size=size)

    A = init_transform(seed, (n, n))
    while not np.isfinite(np.linalg.cond(A)):
        # Keep trying until we get an invertible matrix.
        seed += 1
        A = init_transform(seed, (n, n))

    B = init_transform(seed + 1, (n,))

    def encode(j: int) -> int:
        bits_in = np.array(int_to_bits(j, num_bits(max_i), bool))
        bits_out = np.logical_xor(np.logical_xor.reduce(bits_in * A, axis=1), B)
        #  bits_out = A @ bits_in
        return bits_to_int(bits_out.tolist())

    out = encode(i)
    while out > max_i:
        out = encode(out)

    return out


# Testing
max_i = 10
for i in range(max_i + 1):
    bits = int_to_bits(i, num_bits(max_i))
    assert len(bits) == num_bits(max_i)
    assert i == bits_to_int(bits)

seed = 13
shuffled = set()
for i in range(max_i + 1):
    shuffled_i = shuffled_index(i, max_i, seed)
    shuffled.add(shuffled_i)
    print(i, "->", shuffled_i)
assert len(shuffled) == (max_i + 1)