Fast python algorithm to find all possible partitions from a list of numbers that has subset sums equal to given ratios

Viewed 376

Say I have a list of 20 random integers from 0 to 9. I want to divide the list into N subsets so that the ratio of subset sums equal to given values, and I want to find all possible partitions. I wrote the following code and got it work for the N = 2 case.

import random
import itertools

#lst = [random.randrange(10) for _ in range(20)]
lst = [2, 0, 1, 7, 2, 4, 9, 7, 6, 0, 5, 4, 7, 4, 5, 0, 4, 5, 2, 3]

def partition_sum_with_ratio(numbers, ratios):
    target1 = round(int(sum(numbers) * ratios[0] / (ratios[0] + ratios[1])))
    target2 = sum(numbers) - target1
    p1 = [seq for i in range(len(numbers), 0, -1) for seq in
          itertools.combinations(numbers, i) if sum(seq) == target1
          and sum([s for s in numbers if s not in seq]) == target2]

    p2 = [tuple(n for n in numbers if n not in seq) for seq in p1]

    return list(zip(p1, p2))

partitions = partition_sum_with_ratios(lst, ratios=[4, 3])
print(partitions[0])

Output:

((2, 0, 1, 2, 4, 6, 0, 5, 4, 4, 5, 0, 4, 5, 2), (7, 9, 7, 7, 3))

If you calculate the sum of each subset, you will find the ratio is 44 : 33 = 4 : 3, which are exactly the input values. However, I want the function to work for any number of subsets. For example, I expect

partition_sum_with_ratio(lst, ratios=[4, 3, 3])

to return something like

((2, 0, 1, 2, 4, 6, 0, 5, 4, 4, 3), (5, 0, 4, 5, 2, 7), (9, 7, 7))

I have been thinking about this problem for a month and I found this to be extremely hard. My conclusion is that this problem can only be solved by a recursion. I would like to know if there are any relatively fast algorithm for this. Any suggestions?

1 Answers

Yes, recursion is called for. The basic logic is to do a bipartition into one part and the rest and then recursively split the rest in all possible ways. I've followed your lead in assuming that everything is distinguishable, which creates a lot of possibilities, possibly too many to enumerate. Nevertheless:

import itertools


def totals_from_ratios(sum_numbers, ratios):
    sum_ratios = sum(ratios)
    totals = [(sum_numbers * ratio) // sum_ratios for ratio in ratios]
    residues = [(sum_numbers * ratio) % sum_ratios for ratio in ratios]
    for i in sorted(
        range(len(ratios)), key=lambda i: residues[i] * ratios[i], reverse=True
    )[: sum_numbers - sum(totals)]:
        totals[i] += 1
    return totals


def bipartitions(numbers, total):
    n = len(numbers)
    for k in range(n + 1):
        for combo in itertools.combinations(range(n), k):
            if sum(numbers[i] for i in combo) == total:
                set_combo = set(combo)
                yield sorted(numbers[i] for i in combo), sorted(
                    numbers[i] for i in range(n) if i not in set_combo
                )


def partitions_into_totals(numbers, totals):
    assert totals
    if len(totals) == 1:
        yield [numbers]
    else:
        for first, remaining_numbers in bipartitions(numbers, totals[0]):
            for rest in partitions_into_totals(remaining_numbers, totals[1:]):
                yield [first] + rest


def partitions_into_ratios(numbers, ratios):
    totals = totals_from_ratios(sum(numbers), ratios)
    yield from partitions_into_totals(numbers, totals)


lst = [2, 0, 1, 7, 2, 4, 9, 7, 6, 0, 5, 4, 7, 4, 5, 0, 4, 5, 2, 3]
for part in partitions_into_ratios(lst, [4, 3, 3]):
    print(part)
Related