使用Numba实现线性时间非负整数去重排序时函数失效

核心问题根源 #

Numba的@njit模式对Python原生类型和语法有严格限制，你的函数返回空列表通常是位掩码类型推断错误或位操作溢出导致的——所有元素被误判为已存在，最终没有元素被添加到结果中。

具体排查与修复步骤 #

1. 修复位掩码的类型溢出问题 #

Python原生int是任意精度的，但Numba在@njit模式下会将mask = 0默认推断为固定大小的整数（通常是int32）。如果数组中存在大于31的元素，1 << x会溢出，位掩码无法正确标记元素，所有元素都被判定为已存在，结果列表为空。

修复代码：显式指定掩码为大尺寸无符号整数类型

from numba import njit, uint64

@njit
def unique_linear(arr):
    mask = uint64(0)  # 显式声明64位无符号整数掩码
    res = []
    for x in arr:
        bit = uint64(1) << uint64(x)  # 确保位运算的类型统一
        if not (mask & bit):
            mask |= bit
            res.append(x)
    return res

2. 替换Python列表为NumPy数组 #

Numba对Python列表的append操作支持有限，空列表的类型推断可能出现异常，导致元素无法被正确添加。改用NumPy数组存储结果更可靠。

修复代码：

from numba import njit, uint64
import numpy as np

@njit
def unique_linear(arr):
    if len(arr) == 0:
        return np.array([], dtype=arr.dtype)
    
    # 预先分配结果数组（最大长度为原数组长度）
    res = np.zeros(len(arr), dtype=arr.dtype)
    count = 0
    max_val = arr.max()
    # 处理超过64位的数值，使用多元素掩码
    mask_size = (max_val // 64) + 1
    mask = np.zeros(mask_size, dtype=np.uint64)
    
    for x in arr:
        idx = x // 64
        bit_pos = x % 64
        bit = uint64(1) << uint64(bit_pos)
        if not (mask[idx] & bit):
            mask[idx] |= bit
            res[count] = x
            count += 1
    
    return res[:count]

3. 避免Numba不支持的语法 #

如果你的函数中使用了bin()、in等Python原生操作（比如if x in mask），Numba在@njit模式下无法正确解析这些逻辑，会导致判断条件恒为真或假，最终返回空列表。

修复方式：严格使用位运算(mask & bit) != 0进行存在性判断，不要依赖Python原生的集合或字符串操作。

验证方法 #

在添加@njit装饰器前后分别测试两类用例：

• 小范围数值：[1,2,3,2,1]
• 超出32位的数值：[100, 200, 100]
  对比结果是否一致，确认位掩码类型和位运算逻辑是否正确。

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