CODEWITHSUNDEEP
pythonrecursioninteger-partition
Pobe
Pobe

Reputation: 374

Get n in exact k parts. Recursion and partition algorithm. p(n,k)

I'm looking to enumerate all the partitions of n in k parts.

So for p(5,3) i'd get 2 partitions of k = 3 => (3,1,1), (2,2,1).

Here's what I found from searching and looking through stackoverflow :

def p(n,k):
    lst = []
    if n < k:
        return lst
    elif k == 1:
        return lst
    elif k == n:
        return lst
    else:
        p(n-1, k-1) 
        p(n-k, k)
    return lst

^^^^ This is the form i want,

As it is, finding the sum of k parts is easy, you return p(n-1, k-1) + p(n-k,k). As for me, I need to list each element like so [(3,1,1), (2,2,1)].

My main problem is to "build" those partitions recursively. How would you tackle this?

Edit

If you get base case k = 1, add + 1, k-1 times. (4,1) then (4,1,1)

If you get base case k = n, split up and remove one to each part.

Like so : (3,3) then (3,3,3) then (2,2,2)

If you get base case k < n, nothing

Basically, my problem is to "stack" up the ones from base case to top and get a complete list p(6,3) = [(2,2,2), (4,1,1), (3,2,1)]

enter image description here

Upvotes: 0

Views: 1537

Answers (1)

JuniorCompressor
JuniorCompressor

Reputation: 20025

I would add to the recursive function a third parameter m which is the maximum value an element can have in the partition. Then I would define the function like this:

def p(n, k, m=None):
    if m is None:
        m = n - k + 1 # maximum can be n - k + 1 since minimum is 1
    if k == 0:
        if n == 0:
            yield ()
        return
    for i in xrange(1, m + 1): # first could be from 1 to the maximum
        # the rest of the sum will be n - i among k - 1 elements with
        # maximum i
        for t in p(n - i, k - 1, i):
            yield (i, ) + t

Examples:

>>> list(p(10, 3))
[(4, 3, 3), (4, 4, 2), (5, 3, 2), (5, 4, 1), (6, 2, 2), (6, 3, 1), (7, 2, 1), (8 , 1, 1)]
>>> list(p(6, 2))
[(3, 3), (4, 2), (5, 1)]

Upvotes: 2

Related Questions