Python 3 conversion of a Python 2 implementation of Knuth's "algorithm u" I found online
# Knuth's algorithm to partition ns into m sets. | |
# https://codereview.stackexchange.com/questions/1526/finding-all-k-subset-partitions | |
# http://cs.utsa.edu/~wagner/knuth/ | |
# http://cs.utsa.edu/~wagner/knuth/fasc3b.pdf | |
# ns must be a list of integers | |
# m must be an integer less than len(ns) | |
# this is a generator expression | |
def algorithm_u(ns, m): | |
def visit(n, a): | |
ps = [[] for i in range(m)] | |
for j in range(n): | |
ps[a[j + 1]].append(ns[j]) | |
return ps | |
def f(mu, nu, sigma, n, a): | |
if mu == 2: | |
yield visit(n, a) | |
else: | |
for v in f(mu - 1, nu - 1, (mu + sigma) % 2, n, a): | |
yield v | |
if nu == mu + 1: | |
a[mu] = mu - 1 | |
yield visit(n, a) | |
while a[nu] > 0: | |
a[nu] = a[nu] - 1 | |
yield visit(n, a) | |
elif nu > mu + 1: | |
if (mu + sigma) % 2 == 1: | |
a[nu - 1] = mu - 1 | |
else: | |
a[mu] = mu - 1 | |
if (a[nu] + sigma) % 2 == 1: | |
for v in b(mu, nu - 1, 0, n, a): | |
yield v | |
else: | |
for v in f(mu, nu - 1, 0, n, a): | |
yield v | |
while a[nu] > 0: | |
a[nu] = a[nu] - 1 | |
if (a[nu] + sigma) % 2 == 1: | |
for v in b(mu, nu - 1, 0, n, a): | |
yield v | |
else: | |
for v in f(mu, nu - 1, 0, n, a): | |
yield v | |
def b(mu, nu, sigma, n, a): | |
if nu == mu + 1: | |
while a[nu] < mu - 1: | |
yield visit(n, a) | |
a[nu] = a[nu] + 1 | |
yield visit(n, a) | |
a[mu] = 0 | |
elif nu > mu + 1: | |
if (a[nu] + sigma) % 2 == 1: | |
for v in f(mu, nu - 1, 0, n, a): | |
yield v | |
else: | |
for v in b(mu, nu - 1, 0, n, a): | |
yield v | |
while a[nu] < mu - 1: | |
a[nu] = a[nu] + 1 | |
if (a[nu] + sigma) % 2 == 1: | |
for v in f(mu, nu - 1, 0, n, a): | |
yield v | |
else: | |
for v in b(mu, nu - 1, 0, n, a): | |
yield v | |
if (mu + sigma) % 2 == 1: | |
a[nu - 1] = 0 | |
else: | |
a[mu] = 0 | |
if mu == 2: | |
yield visit(n, a) | |
else: | |
for v in b(mu - 1, nu - 1, (mu + sigma) % 2, n, a): | |
yield v | |
n = len(ns) | |
a = [0] * (n + 1) | |
for j in range(1, m + 1): | |
a[n - m + j] = j - 1 | |
return f(m, n, 0, n, a) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment