使用Numba编译函数时，如何正确比较Numpy字符串数组？

在Numba的nopython=True模式下，直接对Numpy字符串数组使用==比较会返回单个标量布尔值，这是因为Numba未自动将该操作向量化，而是判断整个数组是否完全等于目标字符串。以下是两种可靠的解决方法：

方法一：使用vectorize装饰器实现自动向量化 #

vectorize装饰器会将函数转换为能处理数组元素的向量化函数，自动遍历每个元素完成比较，返回预期的布尔数组：

import numpy as np
from numba import vectorize

test_array = np.array(['1','2','3','4','5'])

@vectorize(nopython=True)
def numba_vectorized(check):
    return check == '1'

# 调用示例
print(numba_vectorized(test_array))
# 输出：[ True False False False False]

方法二：手动遍历数组构建结果 #

如果需要在比较过程中加入额外逻辑，手动遍历数组元素并逐个比较是更灵活的选择：

import numpy as np
from numba import jit

test_array = np.array(['1','2','3','4','5'])

@jit(nopython=True)
def numba_loop_optimised(check):
    # 预先创建与输入数组形状一致的布尔结果数组
    result = np.empty(check.shape, dtype=np.bool_)
    for i in range(check.shape[0]):
        result[i] = check[i] == '1'
    return result

# 调用示例
print(numba_loop_optimised(test_array))
# 输出：[ True False False False False]

原代码问题说明 #

原函数中直接返回check == '1'，在Numba的nopython模式下，该操作被解析为判断整个数组是否所有元素都等于'1'，因此返回单个布尔值而非数组。通过上述两种方法，可实现元素级的比较并得到预期的布尔数组结果。

内容的提问来源于stack exchange，提问作者terrygryffindor