Faster way to perform integer partitioning

Question

Is there a builtin solution (I'd imagine in itertools but couldn't find it myself) to partition an int n (e.g. 4) into its ordered groups of summands (1+1+1+1, 1+1+2, 1+2+1, 2+1+1, ...).

Here's the fastest hand-coded implementation I could find: 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]

Which doesn't generate all permutations, but I can just call set(itertools.permutations on each.

Is there anything faster than that?

Answer 1

You can try to leverage recursion:

def get_comb(target, current=None):
    if current is None:
        current = []

    current_sum = sum(current)

    if current_sum == target:
        yield current
        return

    for n in range(1, target - current_sum + 1):
        yield from get_comb(target, current + [n])


for comb in get_comb(4):
    print(comb)

Prints:

[1, 1, 1, 1]
[1, 1, 2]
[1, 2, 1]
[1, 3]
[2, 1, 1]
[2, 2]
[3, 1]
[4]

---

EDIT: Solution without recursion:

def get_comb(target):
    stack = []
    while True:
        current_sum = sum(stack)
        if current_sum < target:
            stack.append(1)
        elif current_sum > target:
            stack.pop()
        else:
            yield stack
            stack.pop()
            if len(stack) == 0:
                break
            stack.append(stack.pop()+1)

Answer 2

You can think of this problem as n 1s with partitions between each of them, each partition being turned either on or off, so the partitions can be represented by a binary number, and you can simply iterate a number from 0 to 2 ^ (n - 1) - 1 and output 1s while adding a partition (by appending a 0) if the corresponding bit to the iterated number is 1:

def partition(n):
    for i in range(1 << n - 1):
        output = [1]
        for j in range(1, n):
            if 1 << j - 1 & i:
                output.append(0)
            output[-1] += 1
        yield output

print(*partition(4), sep='\n')

This outputs:

[4]
[1, 3]
[2, 2]
[1, 1, 2]
[3, 1]
[1, 2, 1]
[2, 1, 1]
[1, 1, 1, 1]

Demo: https://replit.com/@blhsing/GrownSameClicks

Or if you prefer that the output follow the order in the question, you can iterate backwards:

def partition(n):
    for i in range((1 << n - 1) - 1, -1, -1):
        output = [0]
        for j in range(n, 0, -1):
            if 1 << j - 1 & i:
                output.append(0)
            output[-1] += 1
        yield output

print(*partition(4), sep='\n')

This outputs:

[1, 1, 1, 1]
[1, 1, 2]
[1, 2, 1]
[1, 3]
[2, 1, 1]
[2, 2]
[3, 1]
[4]

Demo: https://replit.com/@blhsing/WholeTurboVoxels

Answer 3

An improved version of Andrej's and a benchmark. Times in milliseconds for n=20, sorted five attempts with each solution:

1119 1286 1288 1299 1368 theonlygusti
1169 1548 1683 1770 1911 Andrej_recursive
 568  620  630  630  764 Andrej_iterative
4281 4491 4539 4569 4732 blhsing
 126  130  147  172  182 Kelly
 185  241  274  302  385 Kelly_copies

My Kelly_copies yields copies of its working list, so all yielded lists are independent. I didn't try the question's idea with itertools.permutations, that's far too slow.

Code (Attempt This Online!):

def Kelly(n):
    stack = [1] * n
    pop = stack.pop
    ones = [None] + [[1] * i for i in range(n)]
    try:
        while True:
            yield stack
            stack[-2] += 1
            stack += ones[pop()]
    except IndexError:
        pass


def Kelly_copies(n):
    stack = [1] * n
    pop = stack.pop
    ones = [None] + [[1] * i for i in range(n)]
    try:
        while True:
            yield stack[:]
            stack[-2] += 1
            stack += ones[pop()]
    except IndexError:
        pass


# copypasted from more-itertools.distinct_permutations, https://more-itertools.readthedocs.io/en/stable/_modules/more_itertools/more.html#distinct_permutations
# this is the special case where items is already sorted and the lengths of the permutations should be the same length as the input
def distinct_permutations(items):
    size = len(items)
    while True:
        # Yield the permutation we have
        yield tuple(items)

        # Find the largest index i such that items[i] < items[i + 1]
        for i in range(size - 2, -1, -1):
            if items[i] < items[i + 1]:
                break
        #  If no such index exists, this permutation is the last one
        else:
            return

        # Find the largest index j greater than j such that items[i] < items[j]
        for j in range(size - 1, i, -1):
            if items[i] < items[j]:
                break

        # Swap the value of items[i] with that of items[j], then reverse the
        # sequence from items[i + 1] to form the new permutation
        items[i], items[j] = items[j], items[i]
        items[i + 1 :] = items[: i - size : -1]  # items[i + 1:][::-1]

# copypasted from https://jeromekelleher.net/generating-integer-partitions.html#most-efficient-algorithm
def accel_asc(n):
    r'''Find all unique permutations of all partitions of 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 from distinct_permutations(a[:k + 2])
            x += 1
            y -= 1
        a[k] = x + y
        y = x + y - 1
        yield from distinct_permutations(a[:k + 1])

def theonlygusti(n):
    return accel_asc(n)


def Andrej_recursive(target, current=None):
    if current is None:
        current = []

    current_sum = sum(current)

    if current_sum == target:
        yield current
        return

    for n in range(1, target - current_sum + 1):
        yield from Andrej_recursive(target, current + [n])


def Andrej_iterative(target):
    stack = []
    while True:
        current_sum = sum(stack)
        if current_sum < target:
            stack.append(1)
        elif current_sum > target:
            stack.pop()
        else:
            yield stack
            stack.pop()
            if len(stack) == 0:
                break
            stack.append(stack.pop()+1)


def blhsing(n):
    for i in range(1 << n - 1):
        output = [1]
        for j in range(1, n):
            if 1 << j - 1 & i:
                output.append(0)
            output[-1] += 1
        yield output


funcs = [
    theonlygusti,
    Andrej_recursive,
    Andrej_iterative,
    blhsing,
    Kelly,
    Kelly_copies,
]

# Correctness
for n in range(1, 15):
    expect = None
    for f in funcs:
        result = sorted(map(tuple, f(n)))
        if expect is None:
            expect = result
        else:
            assert result == expect, (n, f.__name__)

# Speed
from time import time
from collections import deque
ts = {f: [] for f in funcs}
for _ in range(5):
    for f in funcs:
        t0 = time()
        deque(f(20), 0)
        ts[f].append(time() - t0)
        print(*(f'{round(t * 1e3):4}' for t in sorted(ts[f])), f.__name__)
    print()

For the correctness test I skipped n=0, as theonlygusti, Andrej_iterative and blhsing fail that (crash or wrong result) and it's theonlygusti's question, so that failure suggests they're not interested in that case.

Answer 4

A faster but more complicated version of my earlier solution. For n=20 and each solution the best five times of 100 runs (in milliseconds):

  95  110  114  122  125 Kelly
  71   72   79   90   96 Kelly2

My earlier solution Kelly always replaces the current last two numbers so we get the next partition.

The new Kelly2 does the same, but instead uses precomputed lists for how to vary that tail of two numbers. For example if the current tail is [2, 3], then we can replace that with the tails [3, 1, 1], [3, 2], [4, 1] and [5]. They're computed by following([2, 3]) and stored in a lookup table for quick use.

def Kelly(n):
    stack = [1] * n
    pop = stack.pop
    ones = [None] + [[1] * i for i in range(n)]
    try:
        while True:
            yield stack
            stack[-2] += 1
            stack += ones[pop()]
    except IndexError:
        pass


def Kelly2(n):
    ones = [None] + [[1] * i for i in range(n)]
    def following(stack):
        if not all(stack):
            return
        pop = stack.pop
        try:
            while True:
                stack[-2] += 1
                stack += ones[pop()]
                yield stack[:]
        except IndexError:
            pass
    following = [
        [list(islice(following([a, b]), 100))
         for b in range(n - a + 1)]
        for a in range(n)
    ]
    stack = [1] * n
    yield stack
    try:
        while True:
            front, [a, b] = stack[:-2], stack[-2:]
            for back in following[a][b]:
                stack = front + back
                yield stack
    except ValueError:
        pass


from itertools import islice

funcs = [
    Kelly,
    Kelly2,
]

# Correctness
for n in range(1, 15):
    expect = None
    for f in funcs:
        result = sorted(map(tuple, f(n)))
        if expect is None:
            expect = result
        else:
            assert result == expect, (n, f.__name__)

# Speed
from time import time
from collections import deque
ts = {f: [] for f in funcs}
for _ in range(100):
    for f in funcs:
        t0 = time()
        deque(f(20), 0)
        ts[f].append(time() - t0)
for f in funcs:
    print(*(f'{round(t * 1e3):4}' for t in sorted(ts[f])[:5]), f.__name__)

Attempt This Online!

Answer 5

I frankensteined a solution together using the code in the question and more_itertool's distinct_permutations implementation (Prevent duplicates from itertools.permutations):

# copypasted from more-itertools.distinct_permutations, https://more-itertools.readthedocs.io/en/stable/_modules/more_itertools/more.html#distinct_permutations
# this is the special case where items is already sorted and the lengths of the permutations should be the same length as the input
def distinct_permutations(items):
    size = len(items)
    while True:
        # Yield the permutation we have
        yield tuple(items)

        # Find the largest index i such that items[i] < items[i + 1]
        for i in range(size - 2, -1, -1):
            if items[i] < items[i + 1]:
                break
        #  If no such index exists, this permutation is the last one
        else:
            return

        # Find the largest index j greater than j such that items[i] < items[j]
        for j in range(size - 1, i, -1):
            if items[i] < items[j]:
                break

        # Swap the value of items[i] with that of items[j], then reverse the
        # sequence from items[i + 1] to form the new permutation
        items[i], items[j] = items[j], items[i]
        items[i + 1 :] = items[: i - size : -1]  # items[i + 1:][::-1]


# copypasted from https://jeromekelleher.net/generating-integer-partitions.html#most-efficient-algorithm
def accel_asc(n):
    r'''Find all unique permutations of all partitions of 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 from distinct_permutations(a[:k + 2])
            x += 1
            y -= 1
        a[k] = x + y
        y = x + y - 1
        yield from distinct_permutations(a[:k + 1])

This method is about 30% faster than Andrej Kesely's much more readable recursive solution. The accel_asc algorithm is still opaque to me.

>>> timeit.timeit('list(accel_asc(22))', number=3)
7.977816525963135
>>> timeit.timeit('list(get_comb(22))', number=3)
12.344136524014175