| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| from __future__ import annotations |
| from typing import List, Tuple, Optional |
| import random |
|
|
| def stratified_random(n: int=64, seed: Optional[int]=None, shuffle_blocks: bool=True) -> List[int]: |
| if n <= 0 or n & n - 1 != 0: |
| raise ValueError('n must be a positive power of two (e.g., 64)') |
| rng = random.Random(seed) |
| occupied = [[False] * n for _ in range(n)] |
| seq: List[int] = [] |
| blocks: List[Tuple[int, int, int]] = [(0, 0, n)] |
|
|
| def block_has_sample(x0: int, y0: int, size: int) -> bool: |
| for yy in range(y0, y0 + size): |
| row = occupied[yy] |
| for xx in range(x0, x0 + size): |
| if row[xx]: |
| return True |
| return False |
|
|
| def place_random_in_block(x0: int, y0: int, size: int): |
| x = rng.randrange(x0, x0 + size) |
| y = rng.randrange(y0, y0 + size) |
| attempts = 0 |
| while occupied[y][x]: |
| x = rng.randrange(x0, x0 + size) |
| y = rng.randrange(y0, y0 + size) |
| attempts += 1 |
| if attempts > 10000: |
| raise RuntimeError('Too many attempts to place a sample; logic error?') |
| occupied[y][x] = True |
| seq.append(y * n + x) |
| size = n |
| while size > 1: |
| half = size // 2 |
| children: List[Tuple[int, int, int]] = [] |
| for (x0, y0, s) in blocks: |
| assert s == size |
| children.extend([(x0, y0, half), (x0 + half, y0, half), (x0, y0 + half, half), (x0 + half, y0 + half, half)]) |
| if shuffle_blocks: |
| rng.shuffle(children) |
| for (x0, y0, s) in children: |
| if not block_has_sample(x0, y0, s): |
| place_random_in_block(x0, y0, s) |
| blocks = children |
| size = half |
| remaining: List[int] = [] |
| for y in range(n): |
| for x in range(n): |
| if not occupied[y][x]: |
| remaining.append(y * n + x) |
| rng.shuffle(remaining) |
| seq.extend(remaining) |
| assert len(seq) == n * n, (len(seq), n * n) |
| return seq |
|
|
