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)

所以，我希望掩码会是

> [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 ufuncs 的有效实现只针对数字类型给出(我还没有时间查看源代码)。但我不希望更高级别的函数能够直接使用您预先分配的输入数组作为缓冲区。通过编译的 ufunc，我可以想象 out 关键字参数可以让您绕过临时数组的创建，但我认为这里没有其他选择。