python - 将 numpy 数组与标量与 'out=...' 进行比较

Question

我有一个字符串数组:

s = np.array(['a', 'b', 'c'])

我想要一个函数 array_equal_to_scalar 来比较 s 与字符串 'a' 并将输出写入预分配数组(我需要快速性能):

mask = np.empty(s.shape)
np.array_equal_to_scalar(s, 'a', out=mask)

所以，我预计 mask 将会是

> [True False False]

有没有办法制作类似array_equal_to_scalar的东西？

Answer 1

您正在寻找的是 numpy.equal ufunc ，这似乎不适用于您的用例。

为了按照您想要的方式使用它，我们需要显式广播要与适当形状的 numpy 数组进行比较的标量:

import numpy as np

a = np.array(['a','b','c'])
res = np.empty(a.shape, dtype=bool)
np.equal(a, np.broadcast_to(['a'], a.shape), out=res)

不幸的是，上面的调用 (1) 忽略了广播并给出了恒定的结果，并且 (2) 是 NotImplemented。我们可以尝试分配一个适当的比较数组来强制进行适当的元素比较，但没有效果:

>>> compare = np.full(a.shape, 'a')
>>> np.equal(a, compare)
NotImplemented

看来通过 numpy ufunc 的有效实现仅适用于数字类型(我还没有时间研究源代码)。但我不期望更高级别的函数能够直接使用预先分配的输入数组作为缓冲区。通过编译的 ufunc，我可以想象 out 关键字参数可以让您绕过临时数组的创建，但我认为这里没有其他选择。