python - 根据另一个列表或数组查找数组/列表中元素的索引

Question

我有两个列表/数组，如果另一个列表中存在相同的数字，我想找到一个列表中元素的索引。这是一个例子

 list_A = [1,7,9,7,11,1,2,3,6,4,9,0,1]
 list_B = [9,1,7] 
 #output required : [0,1,2,3,5,10,12]

任何使用 numpy 来做到这一点的方法

Answer 1

使用列表理解和 enumerate() :

>>> list_A = [1,7,9,7,11,1,2,3,6,4,9,0,1]
>>> list_B = [9,1,7]
>>> [i for i, x in enumerate(list_A) if x in list_B]
[0, 1, 2, 3, 5, 10, 12]

使用numpy:

>>> import numpy as np
>>> np.where(np.isin(list_A, list_B))
(array([ 0,  1,  2,  3,  5, 10, 12], dtype=int64),)

此外，如@Chris_Rands指出，我们还可以首先将 list_B 转换为集合，因为 in 对于集合来说是 O(1)，而对于列表来说是 O(n)。

时间比较:

import random
import numpy as np
import timeit

list_A = [random.randint(0,100000) for _ in range(100000)]
list_B = [random.randint(0,100000) for _ in range(50000)]

array_A = np.array(A)
array_B = np.array(B)

def lists_enumerate(list_A, list_B):
    return [i for i, x in enumerate(list_A) if x in set(list_B)]

def listB_to_set_enumerate(list_A, list_B):
    set_B = set(list_B)
    return [i for i, x in enumerate(list_A) if x in set_B]

def numpy(array_A, array_B):
    return np.where(np.isin(array_A, array_B))

结果:

>>> %timeit lists_enumerate(list_A, list_B)
48.8 s ± 638 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
>>> %timeit listB_to_set_enumerate(list_A, list_B)
11.2 ms ± 856 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
>>> %timeit numpy(array_A, array_B)
23.3 ms ± 167 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

很明显，对于较大的列表，最佳解决方案是在应用枚举之前将 list_B 转换为集合，或者使用 numpy。