Partitioning an integer with pattern

user_1273987 0
Viewed 89

I am aware that the problem of partitioning an integer is old and there are many questions and answers about it here on SO, but after searching extensively I haven't found exactly what I am looking for. To be fair, my solution is not too too bad, but I'd like to know if there is a faster/better way of doing the following:

I need to partition an integer into a fixed-length partition that may include the value 0, and where each "position" in the partition is subject to a max possible value. For example:

>>>list(partition(number = 5, max_vals = (1,0,3,4)))
[(1, 0, 3, 1),
 (1, 0, 2, 2),
 (1, 0, 0, 4),
 (1, 0, 1, 3),
 (0, 0, 1, 4),
 (0, 0, 2, 3),
 (0, 0, 3, 2)]

My solution is the following:

from collections import Counter
from itertools import combinations

def partition(number:int, max_vals:tuple):
    S = set(combinations((k for i,val in enumerate(max_vals) for k in [i]*val), number))
    for s in S:
        c = Counter(s)
        yield tuple([c[n] for n in range(len(max_vals))])

Essentially I first create "tokens" for each slot, then I combine the right number of them and finally I count how many per slot there are.

I don't particularly like having to instantiate a Counter for each partition, but the thing I dislike the most is that combinations generates many more tuples than what is needed and then I discard all of the duplicates with set(), which seems quite inefficient. Is there a better way?

list algorithm python performance
2 Answers

Even though there must be better algorithms, a relatively simpler and faster solution, using itertools.product will be:

>>> from itertools import product
>>> def partition_2(number:int, max_vals:tuple):
        return (comb for comb in 
                product(*(range(min(number, i) + 1) for i in max_vals)) 
                if sum(comb)==number)

>>> list(partition_2(number = 5, max_vals = (1,0,3,4)))
[(0, 0, 1, 4),
 (0, 0, 2, 3),
 (0, 0, 3, 2),
 (1, 0, 0, 4),
 (1, 0, 1, 3),
 (1, 0, 2, 2),
 (1, 0, 3, 1)]

Performance:

>>> %timeit list(partition(number = 15, max_vals = (1,0,3,4)*3))
155 ms ± 681 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

>>> %timeit list(partition_2(number = 15, max_vals = (1,0,3,4)*3))
14.7 ms ± 763 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
################################################################################
>>> %timeit list(partition(number = 5, max_vals = (10,20,30,10,10)))
1.17 s ± 26.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

>>> %timeit list(partition_2(number = 5, max_vals = (10,20,30,10,10)))
1.21 ms ± 28.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
#################################################################################
>>> %timeit list(partition_2(number = 35, max_vals = (8,9,10,11,12)))
23.2 ms ± 697 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

>>> %timeit list(partition(number = 35, max_vals = (8,9,10,11,12)))
# Will update when/if it finishes :)
user_54317910

A recursive function is usually an elegant way of approaching this kind of problem:

def partition(N,slots):
    if len(slots)==1:
        if slots[0]>=N: yield [N]
        return
    for s in range(min(N,slots[0])+1):
        yield from ([s]+p for p in partition(N-s,slots[1:]))

                    
for part in partition(5,[1,0,3,4]): print(part)
[0, 0, 1, 4]
[0, 0, 2, 3]
[0, 0, 3, 2]
[1, 0, 0, 4]
[1, 0, 1, 3]
[1, 0, 2, 2]
[1, 0, 3, 1]

This can be further optimized by checking the remaining space at each recursion level and short circuit traversal when the remaining slots are insufficient for the number to spread:

def partition(N,slots,space=None):
    if space is None: space = sum(slots)
    if N>space: return
    if len(slots)==1:
        if slots[0]>=N: yield [N]
        return
    for s in range(min(N,slots[0])+1):
        yield from ([s]+p for p in partition(N-s,slots[1:],space-slots[0]))

This optimization improves performance in scenarios where the number of solutions is less than the full product of all slots. It is slower than iterations in cases where most of the slot combinations work.

from timeit import timeit

t = timeit(lambda:list(partition(45,(8,9,10,11,12))),number=1)
print(t) # 0.000679596

t = timeit(lambda:list(partition_2(45,(8,9,10,11,12))),number=1)
print(t) # 0.027492302 (Sayandip's)


t = timeit(lambda:list(partition(15,(1,0,3,4)*3)),number=1)
print(t) # 0.024383259

t = timeit(lambda:list(partition_2(15,(1,0,3,4)*3)),number=1)
print(t) # 0.018362536

To get systematically better performance from the recursive approach, we would need to limit the depth of recursion. This can be done by approaching the problem differently. If we split the slots in two groups and determine the distribution between two combined slots (left and right) we can then apply the partition on each side and combine the results. This will only recurse to a depth of Log2N and will combine large chunks together instead of only adding values one at a time:

from itertools import product
def partition(N,slots,space=None):
    if space is not None and N>space: return
    if len(slots)==1:
        if slots[0]>=N: yield [N]
        return
    if len(slots)==2:
        for left in range(max(0,N-slots[1]),min(N,slots[0])+1):
            yield [left,N-left]
        return
    leftSlots  = slots[:len(slots)//2]
    rightSlots = slots[len(slots)//2:]
    leftSpace,rightSpace = sum(leftSlots),sum(rightSlots)
    for leftN,rightN in partition(N,[leftSpace,rightSpace],leftSpace+rightSpace):
        partLeft  = partition(leftN,  leftSlots,  leftSpace)
        partRight = partition(rightN, rightSlots, rightSpace)
        for leftSide,rightSide in product(partLeft,partRight):
            yield leftSide+rightSide

The performance improvement is then systematic, in all scenarios:

t = timeit(lambda:list(partition(45,(8,9,10,11,12))),number=1)
print(t) # 0.00017742

t = timeit(lambda:list(partition_2(45,(8,9,10,11,12))),number=1)
print(t) # 0.02895038


t = timeit(lambda:list(partition(15,(1,0,3,4)*3)),number=1)
print(t) # 0.00338676

t = timeit(lambda:list(partition_2(15,(1,0,3,4)*3)),number=1)
print(t) # 0.02025453
user_52375600
Related