Elegant Python code for Integer Partitioning
While this answer is fine, I'd recommend skovorodkin's answer.
>>> def partition(number):
... answer = set()
... answer.add((number, ))
... for x in range(1, number):
... for y in partition(number - x):
... answer.add(tuple(sorted((x, ) + y)))
... return answer
...
>>> partition(4)
set([(1, 3), (2, 2), (1, 1, 2), (1, 1, 1, 1), (4,)])
If you want all permutations(ie (1, 3) and (3, 1)) change answer.add(tuple(sorted((x, ) + y))
to answer.add((x, ) + y)
I needed to solve a similar problem, namely the partition of an integer n
into d
nonnegative parts, with permutations. For this, there's a simple recursive solution (see here):
def partition(n, d, depth=0):
if d == depth:
return [[]]
return [
item + [i]
for i in range(n+1)
for item in partition(n-i, d, depth=depth+1)
]
# extend with n-sum(entries)
n = 5
d = 3
lst = [[n-sum(p)] + p for p in partition(n, d-1)]
print(lst)
Output:
[
[5, 0, 0], [4, 1, 0], [3, 2, 0], [2, 3, 0], [1, 4, 0],
[0, 5, 0], [4, 0, 1], [3, 1, 1], [2, 2, 1], [1, 3, 1],
[0, 4, 1], [3, 0, 2], [2, 1, 2], [1, 2, 2], [0, 3, 2],
[2, 0, 3], [1, 1, 3], [0, 2, 3], [1, 0, 4], [0, 1, 4],
[0, 0, 5]
]
A smaller and faster than Nolen's function:
def partitions(n, I=1):
yield (n,)
for i in range(I, n//2 + 1):
for p in partitions(n-i, i):
yield (i,) + p
Let's compare them:
In [10]: %timeit -n 10 r0 = nolen(20)
1.37 s ± 28.7 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [11]: %timeit -n 10 r1 = list(partitions(20))
979 µs ± 82.9 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [13]: sorted(map(sorted, r0)) == sorted(map(sorted, r1))
Out[14]: True
Looks like it's 1370 times faster for n = 20
.
Anyway, it's still far from accel_asc
:
def accel_asc(n):
a = [0 for i in range(n + 1)]
k = 1
y = n - 1
while k != 0:
x = a[k - 1] + 1
k -= 1
while 2 * x <= y:
a[k] = x
y -= x
k += 1
l = k + 1
while x <= y:
a[k] = x
a[l] = y
yield a[:k + 2]
x += 1
y -= 1
a[k] = x + y
y = x + y - 1
yield a[:k + 1]
It's not only slower, but requires much more memory (but apparently is much easier to remember):
In [18]: %timeit -n 5 r2 = list(accel_asc(50))
114 ms ± 1.04 ms per loop (mean ± std. dev. of 7 runs, 5 loops each)
In [19]: %timeit -n 5 r3 = list(partitions(50))
527 ms ± 8.86 ms per loop (mean ± std. dev. of 7 runs, 5 loops each)
In [24]: sorted(map(sorted, r2)) == sorted(map(sorted, r3))
Out[24]: True
You can find other versions on ActiveState: Generator For Integer Partitions (Python Recipe).
I use Python 3.6.1 and IPython 6.0.0.
I've compared the solution with perfplot
(a little project of mine for such purposes) and found that Nolen's top-voted answer is also the slowest.
Both answers supplied by skovorodkin are much faster. (Note the log-scale.)
To to generate the plot:
import perfplot
import collections
def nolen(number):
answer = set()
answer.add((number,))
for x in range(1, number):
for y in nolen(number - x):
answer.add(tuple(sorted((x,) + y)))
return answer
def skovorodkin(n):
return set(skovorodkin_yield(n))
def skovorodkin_yield(n, I=1):
yield (n,)
for i in range(I, n // 2 + 1):
for p in skovorodkin_yield(n - i, i):
yield (i,) + p
def accel_asc(n):
return set(accel_asc_yield(n))
def accel_asc_yield(n):
a = [0 for i in range(n + 1)]
k = 1
y = n - 1
while k != 0:
x = a[k - 1] + 1
k -= 1
while 2 * x <= y:
a[k] = x
y -= x
k += 1
l = k + 1
while x <= y:
a[k] = x
a[l] = y
yield tuple(a[: k + 2])
x += 1
y -= 1
a[k] = x + y
y = x + y - 1
yield tuple(a[: k + 1])
def mct(n):
partitions_of = []
partitions_of.append([()])
partitions_of.append([(1,)])
for num in range(2, n + 1):
ptitions = set()
for i in range(num):
for partition in partitions_of[i]:
ptitions.add(tuple(sorted((num - i,) + partition)))
partitions_of.append(list(ptitions))
return partitions_of[n]
perfplot.show(
setup=lambda n: n,
kernels=[nolen, mct, skovorodkin, accel_asc],
n_range=range(1, 17),
logy=True,
# https://stackoverflow.com/a/7829388/353337
equality_check=lambda a, b: collections.Counter(set(a))
== collections.Counter(set(b)),
xlabel="n",
)