如何改进代码片段，根据条件获取索引？

Question

我想根据条件获取列表中元素的索引。在大多数情况下，可以使用np.where 做到这一点。但是，一旦条件变得更加复杂，事情就会变得一团糟。

就我而言，条件很简单：该元素也存在于第二个列表中。我怎么能写得这么好？我包含了可以工作的示例代码，但由于 for 循环和列表的追加而速度很慢。如果可能的话，我想摆脱两者并拥有高性能的代码。有谁知道如何做到这一点？

import numpy as np

A = np.random.randint(4, size=20)
B = [0, 1, 2]
indices = list()
N = len(A)
for i in range(N):
    if A[i] in B:
        indices.append(i)

Answer 1

这是我能做到的最快速度

[indices.append(i) for i in range(N) if A[i] in B]

这个拍了

Program Executed in 0.000104600000000010

你拍的

Program Executed in 0.000138000000000027