Python implementation of DLX (Algorithm X using dancing links) after Knuth's paper.

dancing_links.py

# Dancing links
# From https://arxiv.org/abs/cs/0011047

from dataclasses import dataclass


@dataclass
class One:
    lt: object
    rt: object
    dn: object
    up: object
    col: object
    row: int

    def __repr__(self):
        return "<One>"


@dataclass
class Col:
    lt: object
    rt: object
    dn: object
    up: object
    col: object
    size: int
    name: str

    def __repr__(self):
        return f"<{self.name}>"

    def __hash__(self):
        return id(self)


@dataclass
class Root:
    lt: object
    rt: object
    rows: int


def new_problem():
    h = Root(None, None, rows=0)
    h.lt = h
    h.rt = h
    return h


def add_column(h, name, optional=False):
    c = Col(None, None, None, None, None, 0, name)
    c.up = c
    c.dn = c
    c.col = c
    if optional:
        c.lt = c
        c.rt = c
    else:
        c.lt = h.lt
        c.rt = h
        h.lt.rt = c
        h.lt = c
    return c


def add_row(h, cols):
    xs = []
    for c in cols:
        xs.append(One(None, None, None, None, col=c, row=h.rows))

    h.rows += 1

    for i, x in enumerate(xs):
        x.lt = xs[i-1]
        x.rt = xs[(i+1) % len(xs)]
        x.up = x.col.up
        x.dn = x.col
        x.col.up = x
        x.up.dn = x


def print_matrix(h):
    rows = {}
    for c in iter_rt(h):
        print(c.name, end=" ")
        for r in iter_dn(c):
            rows[r.row] = set(x.col.name for x in iter_rt(r)) | {c.name}
    print()

    for r in range(h.rows):
        if r not in rows:
            print("--" * len(list(iter_rt(h))))
            continue
        for c in iter_rt(h):
            if c.name in rows[r]:
                print("1", end=" ")
            else: print("0", end=" ")
        print()


def search(h, out=()):
    if h.rt is h:
        yield [[y.col.name for y in iter_rt(x, inclusive=True)] for x in out]
        return
    c = select_column(h)
    cover_column(c)
    for r in iter_dn(c):
        out_ = (*out, r)
        for j in iter_rt(r):
            cover_column(j.col)
        yield from search(h, out_)
        r = out_[-1]
        c = r.col
        for j in iter_lt(r):
            uncover_column(j.col)
    uncover_column(c)


def search_all(h, out=()):
    if h.rt is h:
        return [[[y.col.name for y in iter_rt(x, inclusive=True)] for x in out]]
    c = select_column(h)
    cover_column(c)
    solutions = []
    for r in iter_dn(c):
        out_ = (*out, r)
        for j in iter_rt(r):
            cover_column(j.col)
        solutions.extend(search_all(h, out_))
        r = out_[-1]
        c = r.col
        for j in iter_lt(r):
            uncover_column(j.col)
    uncover_column(c)
    return solutions


def search_one(h, out=()):
    if h.rt is h:
        return [[y.col.name for y in iter_rt(x, inclusive=True)] for x in out]
    c = select_column(h)
    cover_column(c)
    rs = iter_dn(c)
    solution = None
    while True:
        try:
            r = next(rs)
        except StopIteration:
            break

        out_ = (*out, r)
        for j in iter_rt(r):
            cover_column(j.col)
        solution = search_one(h, out_)
        if solution:
            break
        r = out_[-1]
        c = r.col
        for j in iter_lt(r):
            uncover_column(j.col)
    uncover_column(c)
    return solution


def select_column(h):
    # return h.rt  # simple heuristic
    return min(iter_rt(h), key=lambda c: c.size)  # column with minimal branching factor


def cover_column(c):
    assert isinstance(c, Col)
    c.rt.lt = c.lt
    c.lt.rt = c.rt
    for i in iter_dn(c):
        for j in iter_rt(i):
            j.dn.up = j.up
            j.up.dn = j.dn
            j.col.size -= 1


def uncover_column(c):
    assert isinstance(c, Col)
    for i in iter_up(c):
        for j in iter_lt(i):
            j.col.size += 1
            j.dn.up = j
            j.up.dn = j
    c.rt.lt = c
    c.lt.rt = c


def iter_fld(fld):
    def it(c, inclusive=False):
        if inclusive:
            yield c
        i = getattr(c, fld)
        while i is not c:
            yield i
            i = getattr(i, fld)
    return it

iter_dn = iter_fld("dn")
iter_lt = iter_fld("lt")
iter_rt = iter_fld("rt")
iter_up = iter_fld("up")


if __name__ == "__main__":
    # try it...

    prob = new_problem()

    a = add_column(prob, "A")
    b = add_column(prob, "B")
    c = add_column(prob, "C")
    d = add_column(prob, "D")
    e = add_column(prob, "E")
    f = add_column(prob, "F")
    g = add_column(prob, "G")

    #add_row(prob, "ABCDEFG")

    add_row(prob, {c,e,f})
    add_row(prob, {a, d, g})
    add_row(prob, {b,c,f})
    add_row(prob, {a,d})
    add_row(prob, {b,g})
    add_row(prob, {d,e,g})

    print_matrix(prob)

    print("Solutions:")
    for solution in search(prob):
        print(" ", solution)

    print("Solutions:")
    for solution in search_all(prob):
        print(" ", solution)

    print("First Solution:")
    for solution in search_one(prob):
        print(" ", solution)