de_bruijn.py

from typing import List, Sequence, TypeVar, overload

T = TypeVar("T")


@overload
def de_bruijn(alphabet: str, n: int) -> str:
    ...


@overload
def de_bruijn(alphabet: Sequence[T], n: int) -> List[T]:
    ...


def de_bruijn(alphabet: Sequence[T], n: int) -> str | List[T]:
    """
    从 alphabet 生成窗口长度 n 的 de Bruijn 序列
    摘自维基百科，对参数类型和返回类型稍作修改
    """
    size = len(alphabet)
    a = [0] * size * n
    sequence: List[int] = []

    def search(t, p):
        if t > n:
            if n % p == 0:
                sequence.extend(a[1 : p + 1])
            return
        a[t] = a[t - p]
        search(t + 1, p)
        for k in range(a[t - p] + 1, size):
            a[t] = k
            search(t + 1, t)

    search(1, 1)
    if isinstance(alphabet, str):
        return "".join(alphabet[k] for k in sequence)
    return [alphabet[k] for k in sequence]


def pretty_print_binary_de_bruijn_table(bits: int):
    magic = int(de_bruijn("01", bits), 2)
    size = 1 << bits
    mask = size - 1
    bit_perm = [0] * size
    for k in range(size):
        bit_perm[magic << k >> size - bits & mask] = k
    digits = len(str(size))
    print(f"int ctz{size}(uint{size}_t x) {{")
    print(f"    static const int bit_perm[{size}] = {{")
    for l in range(0, size, 16):
        print(
            " " * 7, "".join(f"{index:>{digits}}, " for index in bit_perm[l : l + 16])
        )
    print("    };")
    print(f"    return bit_perm[{magic:#X} * (x & ~x + 1) >> {size - bits}];")
    print("}")


def main():
    pretty_print_binary_de_bruijn_table(7)


if __name__ == "__main__":
    main()