在 python 中将子列表分组到阈值长度下的函数的优雅和优化解决方案

Question

Given a list L , for instance, [[1,1,1,1], [1,1,1], [1,1], [1]] and a max_len=8 I would like to create a new list LN like this [[[1, 1, 1, 1], [1, 1, 1]], [[1, 1], [1]]] .

So I have a list of lists. I want to group the lists in a way that the sum of the lengths of each list is <= max_len. You need to retain the lists as they are, same order, so only consecutive lists can be grouped.

I've been trying to do it in the most Pythonic and efficient way. Should be O(n). With the help of someone, this is the code I have so far:

def chunks(list_to_chunck, max_len):
    if any(len(sub_list) > max_len for sub_list in list_to_chunck):
        return None

    new_list = []
    while list_to_chunck:
        copy_list = [list_to_chunck.pop(0)]
        while list_to_chunck:
            if len(list_to_chunck[0]) + sum(len(sub_list) for sub_list in copy_list) <= max_len:
                copy_list.append(list_to_chunck.pop(0))
            else:
                break
        new_list.append(copy_list)

    return new_list

Answer 1

You can use a variable size to keep track of the current size of the last sub-list in the output list, and whenever it's going to exceed max_len after adding the current sub-list in the iteration, append a new sub-list to the output. Initialize size with a value greater than max_len so that it will always add a new sub-list in the first iteration. With this approach the time complexity will be O(n) :

def chunks(lst, max_len):
    output = []
    size = max_len + 1
    for s in lst:
        size += len(s)
        if size > max_len:
            output.append([])
            size = len(s)
            if size > max_len:
                return
        output[-1].append(s)
    return output

so that chunks([[1, 1, 1, 1], [1, 1, 1], [1, 1], [1]], 8) returns:

[[[1, 1, 1, 1], [1, 1, 1]], [[1, 1], [1]]]

Answer 2

Here is a sketch of an approach that will be O(N). It creates a new list, and doesn't modify the original. It doesn't handle all edge cases, but this should get you going:

In [1]: data = [[1,1,1,1], [1,1,1], [1,1], [1]]
   ...:

In [2]: def chunks(nested, maxlen):
   ...:     total = 0
   ...:     result = []
   ...:     piece = []
   ...:     for sub in nested:
   ...:         length = len(sub)
   ...:         if total + length > maxlen:
   ...:             result.append(piece)
   ...:             piece = [sub]
   ...:             total = length
   ...:         else:
   ...:             piece.append(sub)
   ...:             total += length
   ...:     if piece:
   ...:         result.append(piece)
   ...:     return result
   ...:

In [3]: chunks(data, 8)
Out[3]: [[[1, 1, 1, 1], [1, 1, 1]], [[1, 1], [1]]]

Answer 3

You don't seem to have any constraints on total number of lists, so a greedy approach should be fine:

def chunks(items, max_len):
    ret = [[]]
    remaining = max_len
    for i in items:
        if len(i) > remaining:
            ret.append([])
            remaining = max_len
            if len(i) > remaining:
                return None  # Could raise on impossible
        ret[-1].append(i)
        remaining -= len(i)
    return ret

using your example:

items = [[1,1,1,1], [1,1,1], [1,1], [1]]
assert chunks(items, 8) == [[[1, 1, 1, 1], [1, 1, 1]], [[1, 1], [1]]]

Seeing other answers, this is pretty much in line, so I wanted to toss in a less readable option, without the length guarantee =)

def chunks(items, max_len):
    count = [0, 0]
    def group(item): 
        count[1] += len(item)
        if count[1] >= 8:
            count[0] += 1
            count[1] = 0
        return count[0]
    return [list(v) for k, v in itertools.groupby(data, key=group)]

Answer 4

A generator-based solution:

def group_subseqs(seq, max_len):
    curr_size = 0
    result = []
    for subseq in seq:
        len_subseq = len(subseq)
        if curr_size + len_subseq <= max_len:
            result.append(subseq)
            curr_size += len_subseq
        else:
            if result:
                yield result
                if len_subseq <= max_len:
                    result = [subseq]
                    curr_size = len_subseq
                else:
                    return
            else:
                return
    if result:
        yield result

working like expected (sort of... stops yielding if a sublist is larger than max_len instead of not yielding anything at all):

a = [[1,1,1,1], [1,1,1], [1,1], [1]]
print(list(group_subseqs(a, 8)))
# [[[1, 1, 1, 1], [1, 1, 1]], [[1, 1], [1]]]

print(list(group_subseqs(a, 4)))
# [[[1, 1, 1, 1]], [[1, 1, 1]], [[1, 1], [1]]]

print(list(group_subseqs(a, 3)))
# []

print(list(group_subseqs(a[::-1], 3)))
# [[[1], [1, 1]], [[1, 1, 1]]]