Add items in a list until their sum exceeds a threshold

Viewed 2633

From a list of values I want to create a new list of values until they add up a value.

I am new with Python but I believe it is best done with a while loop.

L = [1,2,3,4,5,6,7,8,9]
i = 0
s = 0
while i < len(L) and s + L[i] < 20:
    s += L[i]
    i += 1
6 Answers

numpy arrays make this simple

import numpy as np

arr = np.array(L)
arr[arr.cumsum() <= 20].tolist()
#[1, 2, 3, 4, 5]

Since you tagged Pandas:

pd.Series(L, index=np.cumsum(L)).loc[:20].values

Output:

array([1, 2, 3, 4, 5], dtype=int64)

You first create the empty list, and then append values with the same conditions you stated. Finally printing the list will return you the values that got added that match your criteria:

L = [1,2,3,4,5,6,7,8,9]
i = 0
s = 0
new_list = []
while i < len(L) and s + L[i] < 20:
    new_list.append(L[i])
    s += L[i]
    i += 1
print(new_list)

Output:

[1, 2, 3, 4, 5]

You can easily write your own generator function:

L = [1,2,3,4,5,6,7,8,9]

def takeuntil(lst, max_value = 0):
    """Yields elements as long as the cummulative sum is smaller than max_value."""
    total = 0
    for item in lst:
        total += item
        if total <= max_value:
            yield item
        else:
            raise StopIteration

    raise StopIteration

new_lst = [item for item in takeuntil(L, 20)]
print(new_lst)

Which yields

[1, 2, 3, 4, 5]

If you would sum from start to end the (imo) more pythonic way would be

s = 0
for element in L:
   if s + element < 20:
         s += element
    else:
         break

If we’re talking pythonic, a for loop makes more sense (also using better variable names):

data = [1,2,3,4,5,6,7,8,9]
filtered = []
for num in data:
    if num < 20:
        filtered.append(num)

But a comprehension is also pythonic and shorter:

filtered = [num for num in data if num < 20]

Then to get the sum just use the sum function:

total = sum(filtered)

Or if you only need the sum:

total = sum(n for n in data if n < 20)
Related