Python在列表a中查找列表b中项的索引

Question

我有两个列表，如下所示：

aa=[int(1000*random.random()) for i in xrange(10000)]
bb=[int(1000*random.random()) for i in xrange(10000)]

我希望能够得到另一个列表，告诉我在列表bb中项aa的位置；如果它不存在，则返回-1。

这些列表很可能非常庞大，因此需要运行成千上万次，因此即使时间加速也会非常巨大。

到目前为止，我找到的最快的方法是：

def index_withoutexception(aa,bb):
    try:
        return aa.index(bb)
    except:
        return -1
ls = [index_withoutexception(bb,i) for i in aa]

有没有更快的方法来实现这个？

注意：if语句存在问题，因为我找不到一个返回nan/-1的函数，它们都会抛出异常，这是速度慢的原因...我收集到的信息。

- bpython

---

使用字典而不是列表？这样查找时间将会是O(1)，而不是O(n)。 - Tadhg McDonald-Jensen

就速度和Python的风格而言，您的方法没有任何问题。 - muratgu

如果aa中的项目在bb中列出两次怎么办？index()只会返回第一次出现的位置。你想如何处理？ - user6025378

我可以确认这些列表将是唯一的，我只是为了问题而使用了一个随机数数组。在这个例子中，为了方便起见，对于多个条目，第一个条目就足够了。 - bpython

Answer 1

numpy_indexed 软件包可以完全向量化地解决这个问题（免责声明：我是它的作者）。请注意，您最好也用numpy替换掉其余的代码，否则那将成为瓶颈。

import numpy_indexed as npi
i = npi.indices(aa, bb, missing='mask').filled(-1)

- Eelco Hoogendoorn

Answer 2

这是一种基于np.searchsorted的方法，受到此帖子的启发。

sidx = np.argsort(bb)
L = np.searchsorted(bb,aa,sorter=sidx,side='left')
R = np.searchsorted(bb,aa,sorter=sidx,side='right')
out = np.where(L != R,sidx[L],-1)

请注意，如果bb已经排序，您可以跳过计算sidx和所有其他带有sidx的部分，从而提高性能。这种情况下的缩短代码如下：-

L = np.searchsorted(bb,aa,side='left')
R = np.searchsorted(bb,aa,side='right')
out = np.where(L != R,L,-1)

请注意，输出将是一个NumPy数组。如果需要列表输出，可以使用out.tolist()。

运行时测试

让我们对所提出的方法进行计时，并与原始循环版本进行比较。

1] 设置输入：

In [171]: import numpy as np
     ...: 
     ...: # Create random unique lists
     ...: 
     ...: # 1. Random elements
     ...: aa=[int(1000*np.random.random()) for i in xrange(10000)]
     ...: bb=[int(1000*np.random.random()) for i in xrange(10000)]
     ...: 
     ...: # 2. Unique elements
     ...: aa = np.unique(aa)
     ...: bb = np.unique(bb)
     ...: 
     ...: # 3. Since np.unique sorts the elements, let's randomize them
     ...: aa = aa[np.random.permutation(aa.size)]
     ...: bb = bb[np.random.permutation(bb.size)]
     ...: 
     ...: #4. Finall make lists from the arrays
     ...: aa = aa.tolist()
     ...: bb = bb.tolist()
     ...: 

2] 定义循环和向量化版本：

In [172]: def index_withoutexception(aa,bb):
     ...:     try:
     ...:         return aa.index(bb)
     ...:     except:
     ...:         return -1
     ...:     

In [173]: def vectorized_approach(aa,bb):
     ...:     sidx = np.argsort(bb)
     ...:     L = np.searchsorted(bb,aa,sorter=sidx,side='left')
     ...:     R = np.searchsorted(bb,aa,sorter=sidx,side='right')
     ...:     return np.where(L != R,sidx[L],-1)
     ...: 

3] 最后验证并计时结果：

In [174]: out1 = [index_withoutexception(bb,i) for i in aa]

In [175]: out2 = vectorized_approach(aa,bb)

In [176]: np.allclose(out1,out2)
Out[176]: True

In [177]: %timeit [index_withoutexception(bb,i) for i in aa]
100 loops, best of 3: 11.6 ms per loop

In [178]: %timeit vectorized_approach(aa,bb)
1000 loops, best of 3: 780 µs per loop

- Divakar

---

虽然这很快，但它正在对列表进行排序，而我需要保持列表的原始顺序并根据此获取它们的索引。 - bpython

如果您在谈论使用sidx = np.argsort(bb)的第一个版本，它会使用这些argsort-ed indices根据原始顺序获取它们的索引。在验证结果部分，我们执行np.allclose(out1,out2)来确认这一点，这将输出True，表示向量化和原始版本之间的结果匹配。或者我在这里误解了什么？ - Divakar

抱歉有些含糊，但我无法重新排序列表。在发布问题时，我忘记了np.unique会给出一个有序列表。我正在使用的列表是按特定顺序排列的，因此需要保持该顺序。因此，我正在使用for循环并在我的示例中捕获异常。 - bpython

@bpython那就有问题了。你能把它转换成一个普通的有序列表，或者更好地转换成NumPy数组，然后再使用这种方法吗？ - Divakar

Answer 3

你可以创建一个字典或defaultdict(list)，将每个元素映射到它出现的索引（或索引）。这样，你需要一些更多的空间（比原始列表略多，但仍在同一范围内），但一旦创建了字典，每个索引查找将是O(1)。

>>> lst = [random.randint(0, 100) for _ in range(100)]
>>> indices = collections.defaultdict(list)
>>> for i, e in enumerate(lst):
...     indices[e].append(i)
...
>>> indices[30]
[21, 28, 89]

针对您的具体问题，您可以尝试类似于以下的方法：

>>> aa = [random.randint(0, 10) for _ in range(20)] # [3, 9, 4, 5, 6, 5, 2, 4, 7, 4, 4, 9, 10, 8, 8, 7, 6, 3, 3, 3]
>>> bb = [random.randint(0, 10) for _ in range(20)] # [10, 7, 4, 9, 8, 4, 10, 7, 9, 1, 4, 8, 8, 3, 8, 0, 1, 10, 1, 6]
>>> aa_indices = {e: i for (i, e) in reversed(list(enumerate(aa)))} # {2: 6, 3: 0, 4: 2, 5: 3, 6: 4, 7: 8, 8: 13, 9: 1, 10: 12}
>>> b_in_a = [aa_indices.get(b, -1) for b in bb]
>>> b_in_a
[12, 8, 2, 1, 13, 2, 12, 8, 1, -1, 2, 13, 13, 0, 13, -1, -1, 12, -1, 4]

注意：这里使用 reversed，因为否则字典会包含给定元素的 最后一个 索引。

使用 IPython 的 %timeit 进行一些时间分析：这种方法仅需要 2.24 毫秒来创建字典和另外 2.88 毫秒来创建最终列表，而原始方法需要 173 毫秒。

>>> %timeit [index_withoutexception(bb,i) for i in aa]
10 loops, best of 3: 173 ms per loop
>>> %timeit bb_indices = {e: i for (i, e) in reversed(list(enumerate(bb)))}
100 loops, best of 3: 2.24 ms per loop
>>> %timeit [bb_indices.get(i, -1) for i in aa]
100 loops, best of 3: 2.88 ms per loop

- tobias_k