从列表中获取满足list[i][0] == 'value'的前n个元素

Question

我怎么才能从一个列表的列表中获取前 n 个元素，条件是 list[i][0] == 'value' 呢？

目前我做了这两个步骤：

items = [['value', 0], ['foo', 1], ['value', 2], ['value', 3], ['bar', 4]]
n = 2

tmp_list = []
for item in items:
    if item[0] == 'value':
        tmp_list.append(item)

i = 0
while i < n:
    print(tmp_list[i])
    i += 1

在这个例子中，我期待的结果是 ['value', 0] 和 ['value', 2]。

这个方法可以用，但看起来不太整洁。我总觉得还有更好的办法，只是找不到。

Answer 1

这里有三种解决这个问题的方法。

懒惰的方法

使用一个返回列表的lambda表达式，然后再进行切片操作。

take = lambda arr, key, n: [e for e in arr if e[0] == key][:n]

if __name__ == '__main__':
    arr = [['value', 0], ['foo', 1], ['value', 2], ['value', 3], ['bar', 4]]
    print(take(arr, 'value', 2))

更好的方法

创建一个函数，当达到想要的数量时就停止迭代。

def take(arr, key, amount):
    res = []
    for item in arr:
        if item[0] == key and amount > 0:
            res.append(item)
            amount -= 1
    return res

if __name__ == '__main__':
    arr = [['value', 0], ['foo', 1], ['value', 2], ['value', 3], ['bar', 4]]
    print(take(arr, 'value', 2))

最佳的方法

通过使用生成器和islice来分离amount，这样可以获取到想要的数量。

from itertools import islice

def take(arr, key):
    for item in arr:
        if item[0] == key:
            yield item

if __name__ == '__main__':
    arr = [['value', 0], ['foo', 1], ['value', 2], ['value', 3], ['bar', 4]]
    print(list(islice(take(arr, 'value'), 2)))

生成器的简写形式是下面的列表推导：

def take(arr, key):
    return (item for item in arr if item[0] == key)

Answer 2

第一个改进是可以用列表推导式来简化第一步，用切片来处理第二步：

tmp_list = [item for item in items if item[0] == 'value']

for val in tmp_list[:n]:
    print(val)

然后你可以把它改成一个生成器，这样可以延迟计算。是否算是改进还有争议，这要看你打算如何使用这些值。

from itertools import islice

values = islice((item for item in items if item[0] == 'value'), n)

for val in values:
    print(val)