Last active
February 16, 2023 02:44
-
-
Save epwalsh/e3bee2560cf96bcc3083f15e23dbc030 to your computer and use it in GitHub Desktop.
Quick and easy way convert an array index into an index into a permutation of the array
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment