Partitions of n into k parts with restrictions

Question

I need an algorithm that produces a partition of the number n into k parts with the added restrictions that each element of the partition must be between a and b. Ideally, all possible partitions satisfying the restrictions should be equally likely. Partitions are considered the same if they have the same elements in different order.

For example, with n=10, k=3, a=2, b=4 one has only {4,4,2} and {4,3,3} as possible outcomes.

Is there a standard algorithm for such a problem? One can assume that at least one partition satisfying the restrictions always exists.

I answered a very similar question recently, with an algorithm that computes the conditional probabilities and samples accordingly: http://stackoverflow.com/q/41871179/2144669 — David Eisenstat, Feb 01 '17 at 15:30

tobias_k · Accepted Answer · 2017-02-01T15:24:21.357

You can implement this as a recursive algorithm. Basically, the recurrence is like this:

if k == 1 and a <= n <= b, then the only partition is [n], otherwise none
otherwise, combine all the elements x from a to b with all the partitions for n-x, k-1
to prevent duplicates, also substitute the lower bound a with x

Here's some Python (aka executable pseudo-code):

def partitions(n, k, a, b):
    if k == 1 and a <= n <= b:
        yield [n]
    elif n > 0 and k > 0:
        for x in range(a, b+1):
            for p in partitions(n-x, k-1, x, b):
                yield [x] + p

print(list(partitions(10, 3, 2, 4)))
# [[2, 4, 4], [3, 3, 4]]

This could be further improved by checking (k-1)*a and (k-1)*b for the lower and upper bounds for the remaining elements, respectively, and restricting the range for x accordingly:

        min_x = max(a, n - (k-1) * b)
        max_x = min(b, n - (k-1) * a)
        for x in range(min_x, max_x+1):

For partitions(110, 12, 3, 12) with 3,157 solutions, this reduces the number of recursive calls from 638,679 down to 24,135.

Hm, guess I missed the part that you want a single, random partition, and not the entire list... well, you could still generate them all and then pick one randomly. This would also guarantee that each legal permutation is equally likely. — tobias_k, Feb 01 '17 at 15:40

jdehesa · Answer 2 · 2017-02-01T15:17:35.047

If k and b - a are not too big you can try a randomized depth-first search:

import random

def restricted_partition_rec(n, k, min, max):
    if k <= 0 or n < min:
        return []
    ps = list(range(min, max + 1))
    random.shuffle(ps)
    for p in ps:
        if p > n:
            continue
        elif p < n:
            subp = restricted_partition(n - p, k - 1, min, max)
            if subp:
                return [p] + subp
        elif k == 1:
            return [p]
    return []

def restricted_partition(n, k, min, max):
    return sorted(restricted_partition_rec(n, k, min, max), reverse=True)

print(restricted_partition(10, 3, 2, 4))
>>>
[4, 4, 2]

Although I'm not sure if all the partitions have exactly the same probability in this case.

The values I am interested are n from 38 to 110, k from 6 to 12, a=3, b=12 — Jesús Ros, Feb 01 '17 at 15:08

score 1 · Answer 3 · answered Feb 01 '17 at 23:16

Here's a sampling algorithm that uses conditional probability.

import collections
import random

countmemo = {}


def count(n, k, a, b):
    assert n >= 0
    assert k >= 0
    assert a >= 0
    assert b >= 0
    if k == 0:
        return 1 if n == 0 else 0
    key = (n, k, a, b)
    if key not in countmemo:
        countmemo[key] = sum(
            count(n - c, k - 1, a, c) for c in range(a, min(n, b) + 1))
    return countmemo[key]


def sample(n, k, a, b):
    partition = []
    x = random.randrange(count(n, k, a, b))
    while k > 0:
        for c in range(a, min(n, b) + 1):
            y = count(n - c, k - 1, a, c)
            if x < y:
                partition.append(c)
                n -= c
                k -= 1
                b = c
                break
            x -= y
        else:
            assert False
    return partition


def test():
    print(collections.Counter(
        tuple(sample(20, 6, 2, 5)) for i in range(10000)))


if __name__ == '__main__':
    test()

Partitions of n into k parts with restrictions

3 Answers3