python - 对 numpy 数组有多个相等条件

Question

假设我有一个大型二维 numpy 数组(称之为 A)，其中包含从 0 到 9 的整数。

我正在尝试编写一个函数，该函数返回一个与具有以下特征的 A 形状相同的二进制 numpy 数组(称为 B)。

如果 A 中的相应元素出现在给定列表 L 中，则 B 中的条目为 1；否则，条目为零。

下面显示了执行此操作可能不是最有效的代码。

A = np.random.randint(0, 10, (5, 5))
L = [3,4,5]

B = np.zeros(A.shape)
for e in L:
    B[A==e] = 1

有没有更快的方法？

谢谢!

Answer 1

看来我有责任指出显而易见的事实:

def AinL(A, L):
    B = np.zeros((10,), int)
    B[L] = 1
    return B[A]

基准:

 10x10 #L=3
orig     0.6665631101932377
HAL      0.4370500799268484
Psidom   1.13961720908992
PP       0.23527960386127234

 100x100 #L=3
orig     0.3015591569710523
HAL      0.29902734607458115
Psidom   0.4470538650639355
PP       0.18963343487121165

 1000x1000 #L=4
orig     0.5516874771565199
HAL      0.5967503408901393
Psidom   0.6331975681241602
PP       0.23225238709710538

 10000x1000 #L=2
orig     0.8539429588709027
HAL      0.9840140701271594
Psidom   1.0392512339167297
PP       0.7203555379528552