chunk list of objects by property sum

user_1914034 0
Viewed 119

I want to chunk a list of objects in order that the sum of a property is less than a value, let say 10.

lst = [{"value": 1}, {"value": 8}, {"value": 2}, {"value": 6}, {"value": 1}, {"value": 9}, {"value": 10}]

result = [[{"value": 1}, {"value": 8}], [{"value": 2}, {"value": 6}, {"value": 1}], [{"value": 9}], [{"value": 10}]]

This is what I tried so far:

def chunks_by_prop_sum(lst, prop, max_per_chunk):
    chunks = []
    chunk_prop_sum = 0
    chunk = []
    for i in range(0, len(lst)):
        if chunk_prop_sum <= max_per_chunk- lst[i][prop]
            chunk.append(lst[i])
            chunk_prop_sum += lst[i][prop]
        else:
            chunks.append(chunk.copy())
            chunk = []
            chunk_prop_sum = 0

    return chunks

I would like to improve that function by making a generator if possible and also to handle when an object value is greater than the max per chunk threshold.

The order of objects in chunks matter.

python
3 Answers

This is building a generator, but also without using any if statements and not a lot of code.

from collections import defaultdict

def chunks_by_prop_sum(data:list[dict], prop:str, max_per_chunk:int) -> list[list[dict]]:
    res = defaultdict(list)
    total, inc = 0, 0
    for entry in data:
        total += entry[prop]
        res[(total // max_per_chunk) + inc].append(entry)
        inc += int(entry[prop] >=  max_per_chunk)
    return list(res.values())

print(chunks_by_prop_sum(lst, 'value', 10))
user_140787580

Here's one way of converting your function to a generator:

original = [{"value": 1}, {"value": 8}, {"value": 2}, {
    "value": 6}, {"value": 1}, {"value": 9}, {"value": 10}]


def my_chunks(lst, val, n):
    chunk, chunk_sum = [], 0
    while chunk_sum < n and lst:
        if lst[0][val] >= n:
            if chunk:
                yield chunk
            yield [lst.pop(0)]
            chunk, chunk_sum = [], 0
        # We could also use the walrus operator here
        # elif (tot:= chunk_sum + lst[0][val]) < n:
        #   chunk_sum = tot
        elif chunk_sum + lst[0][val] < n:
            chunk_sum += lst[0][val]
            chunk.append(lst.pop(0))
        else:
            yield chunk
            chunk, chunk_sum = [], 0
    if chunk:
        yield chunk


result = [[{"value": 1}, {"value": 8}], [{"value": 2}, {
    "value": 6}, {"value": 1}], [{"value": 9}], [{"value": 10}]]
print([*my_chunks(original, 'value', 10)] == result)

Outputs:

True

The idea here is to build a chunk in the while loop and yield then "reset" it when the sum of the objects in the chunk is greater than n. When the loop encounters an object with a value greater than n it yields the prior chunk if it exists then the object in a chunk of its own.

user_164501690

There is another option. It is quite neat:

def chunks_by_prop_sum(lst, prop, max_per_chunk):
    out = []
    cum = 0
    for item in lst:
        if out and cum + item[prop] > max_per_chunk:
            yield out
            out = []
            cum = 0
        out.append(item)
        cum += item[prop]
    yield out
    
rv = list(chunks_by_prop_sum(lst, "value", 10))
print(rv == result)

Output:

True
user_174570420
Related