Generate Necklace, Lyndon, and de Bruijn sequences

seq.py

#!/usr/bin/env python3

import argparse

from typing import Callable, Generator, TypeVar


def main():
    parser = argparse.ArgumentParser(
        "Generates necklace-related sequences using the FKM algorithm as described by \"Combinatorial Generation\" by Frank Ruskey. See https://web.archive.org/web/20221203232527/https%3A%2F%2Fpage.math.tu-berlin.de%2F~felsner%2FSemWS17-18%2FRuskey-Comb-Gen.pdf"
    )
    parser.add_argument("-k", type=int, help="Alphabet size", required=True)
    parser.add_argument("-n", type=int, help="Code length", required=True)
    parser.add_argument(
        "-i",
        type=bool,
        action=argparse.BooleanOptionalAction,
        help=
        "Interactive mode. Stdin must issue a newline after each sequence item."
    )
    parser.add_argument(
        "-c",
        type=bool,
        action=argparse.BooleanOptionalAction,
        help=
        "Complete the sequence. This appends the first 'n-1' characters to the sequence, meaning the output will contain all unique n-length strings of the alphabet without needing to cycle."
    )
    args = parser.parse_args()
    # TODO: Choose sequence type via subcommand.
    sequence = fkm_recursive(args.k, args.n, extract_de_bruijn)
    # sequence = fkm_recursive(args.k, args.n, extract_pre_necklaces)
    # sequence = fkm_recursive(args.k, args.n, extract_lyndon)
    # sequence = fkm_recursive(args.k, args.n, extract_necklaces)
    if args.c:
        # Save and replay the first n-1 items.
        sequence = replay(sequence, args.n - 1)
    if args.i:
        for code in sequence:
            print("".join(str(c) for c in code), end="")
            input("")
    else:
        for code in sequence:
            print("".join(str(c) for c in code))


def extract_pre_necklaces(n: int, p: int,
                          a: list[int]) -> Generator[list[int], None, None]:
    # NOTE: This breaks with n > 10. Consider allowing different alphabets in
    # that case (and below).
    yield [x for x in a[1:n + 1]]


def extract_lyndon(n: int, p: int,
                   a: list[int]) -> Generator[list[int], None, None]:
    if p == n:
        yield [x for x in a[1:n + 1]]


def extract_necklaces(n: int, p: int,
                      a: list[int]) -> Generator[list[int], None, None]:
    if n % p == 0:
        yield [x for x in a[1:n + 1]]


def extract_de_bruijn(n: int, p: int,
                      a: list[int]) -> Generator[list[int], None, None]:
    if n % p == 0:
        yield from ([x] for x in a[1:p + 1])


def fkm_recursive(
    k: int, n: int, extract_seq: Callable[[int, int, list[int]],
                                          Generator[list[int], None, None]]
) -> Generator[list[int], None, None]:
    a = [0] * k * n

    def fkm(t: int, p: int) -> Generator[list[int], None, None]:
        if t > n:
            yield from extract_seq(n, p, a)
        else:
            a[t] = a[t - p]
            yield from fkm(t + 1, p)
            for j in range(a[t - p] + 1, k):
                a[t] = j
                yield from fkm(t + 1, t)

    yield from fkm(1, 1)


T = TypeVar("T")


# Replay the first n items from the given generator at the end of the given
# generator after yielding all values. If len(g) <= n, then the entire contents
# of g will be replayed.
def replay(g: Generator[T, None, None], n: int) -> Generator[T, None, None]:
    extension: list[T] = []
    for value in g:
        if len(extension) < n:
            extension.append(value)
        yield value
    yield from extension


if __name__ == "__main__":
    main()