python - 切割列表列表

Question

给定一个限制和包含任意数量项目的列表列表，我想削减列表列表，其中限制是所有列表中项目的总数。

我目前的努力:

lol = [['a', 'v', '1', 'r'], ['d', 'x', 'b'], ['s', '0', 'u', 'x']]

def cut_lol(lol, n):
    count = 0
    result = []
    for l in lol:
        if count + len(l) > (n):
            new_len_l = count + len(l) - n
            result.append(l[:new_len_l])
            return result
        result.append(l)
        count += len(l)
    return result

但是这个函数有问题，我看不出是什么。

例如:

print(cut_lol(lol, 8))
# returns    [['a', 'v', '1', 'r'], ['d', 'x', 'b'], ['s', '0', 'u']]
# should be: [['a', 'v', '1', 'r'], ['d', 'x', 'b'], ['s']]


print(cut_lol(lol, 4))
# returns    [['a', 'v', '1', 'r'], ['d', 'x', 'b']]
# should be: [['a', 'v', '1', 'r']]


print(cut_lol(lol, 3))
# returns    [['a']]
# should be: [['a', 'v', '1']]

谁能帮我完成这项工作？谢谢!

Answer 1

使其工作的最小更改:

lol = [['a', 'v', '1', 'r'], ['d', 'x', 'b'], ['s', '0', 'u', 'x']]

def cut_lol(lol, n):
    count = 0
    result = []
    for l in lol:
        if count + len(l) > (n):
            new_len_l = n - count
            if new_len_l:
                result.append(l[:new_len_l])
            return result
        result.append(l)
        count += len(l)
    return result

简化版:

lol = [['a', 'v', '1', 'r'], ['d', 'x', 'b'], ['s', '0', 'u', 'x']]

def cut_lol(lol, n):
    count = 0
    result = []
    for l in lol:
        to_read = n - count
        if to_read <= 0:
            break
        result.append(l[:to_read])
        count += len(l)
    return result