python - 列表理解索引

Question

例如，假设您有列表

a = [['hand', 'head'], ['phone', 'wallet'], ['lost', 'stock']]

和另一个清单

b = ['phone', 'lost']

并且您想找到一个 list c，其中包含 in 中的行的索引a（将a其视为 2D 矩阵），其第一列是 in 中的值b。所以在这种情况下

c = [1, 2]

我尝试使用以下列表推导

c = [i if a[i][0] in b for i in range(0, 1)]
c = [i if a[i][0] in b]

但是这两个都是无效的语法。

Answer 1

使用enumerate()：

c = [i for i, v in enumerate(a) if v[0] in b]

enumerate()为您提供传入的可迭代对象的索引和值。请注意，if测试在最后进行；列表推导式的编写顺序应与嵌套循环时使用的顺序相同：

c = []
for i, v in enumerate(a):
    if v[0] in b:
        c.append(i)  

你真的很想做b一套：

b = set(b)

使成员资格测试成为 O(1) 恒定时间操作，而不是针对列表的 O(n) 线性时间测试。

演示：

>>> a = [['hand', 'head'], ['phone', 'wallet'], ['lost', 'stock']]
>>> b = {'phone', 'lost'}  # set literal
>>> [i for i, v in enumerate(a) if v[0] in b]
[1, 2]

Answer 2

您也可以使用 Numpy 来执行此操作：

>>> import numpy as np
>>> a = np.array([['hand', 'head'], ['phone', 'wallet'], ['lost', 'stock']])
>>> b = np.array(['phone', 'lost'])
>>> np.in1d(a[:,0],b)
array([False,  True,  True], dtype=bool)

或者，如果您想要索引：

>>> np.where(np.in1d(a[:,0],b))[0].tolist()
[1, 2]

Answer 3

首先数组从 0 开始。所以 c 必须是：

c=[1,2]

如果您需要使用列表压缩来执行此操作，则解决方案可以是：

c=[pos for pos, val_a in enumerate(a) for val_b_to_check in val_a if val_b_to_check in b ]