简体   繁体   中英

python get a list of length X where the sum of all values is Y

I need to find all the combinations of lists who has their sum equal to X in the shortest possible time.

At this moment i have this:

def deduceArrayFromSum(arr, expectedSum, depth, maxDepth, intervalEnd, intervalStart):
    if maxDepth<=maxOffset:
        if expectedSum>0:
            i = min(intervalEnd, asum)
            while i>=intervalStart:
                if expectedSum>=i and depth<maxOffset:
                    arr[depth] = i
                    deduceArrayFromSum(arr, expectedSum-i, depth+1, maxDepth, i, intervalStart)
                i=i-1
        elif asum==0:
            foundSum(arr)

For every combination found, it calls foundSum()

I would expect that i could optimize this by making it linear instead of recursive calls, and by avoiding sending the arr variable at every call. Another possible way to speed this up would be by using yield , but I cannot seem to understand how it works.

Maybe numpy can help somehow?


Edit: arr starts as [0, 0, 0] when calling deduceArrayFromSum(arr, 100, 0, len(arr), 100, 0) , foundSum() gets called with the following params:

[100, 0, 0]
[99, 1, 0]
[98, 2, 0]
[98, 1, 1]
[...]

I hope this clears what i want from this code

Use itertools.combinations .

import itertools

def deduceArrayFromSum(arr, length, expectedSum):
    for combination in itertools.combinations(arr, length):
        if sum(combination) == expectedSum:
            foundSum(combination)

With a little refactoring, I would turn it into a generator:

import itertools

def combinations_with_sum(iterable, length, sum_):
    return (c in itertools.combinations(iterable, length) if sum(c) == sum_)

Then you could use it like this:

for t in combinations_with_sum([1, 2, 3, 4], 2, 5):
    print t

or

print list(combinations_with_sum([1, 2, 3, 4], 2, 5))

This problem is NP-complete , so it will always be slow for large inputs. But you may consider researching better algorithms than brute force.

You can do this:

>>> from itertools import combinations_with_replacement as it_cr
>>> len_=3
>>> tgt=100
>>> li=[t[::-1] for t in it_cr(range(tgt+1), len_) if sum(t)==tgt]
>>> li
[(100, 0, 0), (99, 1, 0), (98, 2, 0), (97, 3, 0), (96, 4, 0), (95, 5, 0),
...
 (35, 33, 32), (34, 34, 32), (34, 33, 33)]

Which is easily turned into a generator:

def tups_sum_x(tgt, len_):
    for t in it_cr(range(tgt+1), len_):
        if sum(t)==tgt:
            yield t[::-1]

If you want the same order that you present, just sort:

>>> sorted(li, key=lambda t: (-t[0], t[1:]))
[(100, 0, 0), (99, 1, 0), (98, 1, 1), (98, 2, 0), (97, 2, 1), (97, 3, 0), ...

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM