使用Numba优化位图排序去重：深入理解整数表示与潜在陷阱

import numpy as np
from time import perf_counter
from numba import njit

def count_unique_bitmask(ls):
    """
    使用位图法对非负整数进行排序并去重。
    该方法假设输入整数非负且在位图m所能表示的范围内。
    """
    ret = []
    m = 0  # 初始化位图，用于标记数字的存在

    # 阶段1: 构建位图
    for x in ls:
        # 将数字x对应的位设置为1
        # 例如，如果x=3，则m的第3位变为1
        m = m | (1 << int(x))

    # 阶段2: 从位图中提取排序后的唯一元素
    i = 0
    while m > 0: # 只要位图m中还有位是1，就继续
        if (m & 1): # 检查当前最低位（第i位）是否为1
            ret.append(i)
        m = m >> 1 # 位图右移一位，准备检查下一位
        i += 1     # 对应数字递增
    return ret

# 示例测试
RNG = np.random.default_rng(0)
# 生成一个包含大量随机非负整数的数组，最大值2^16-1，数量2^17
x = RNG.integers(2**16, size=2**17) 
print(f"输入数组大小: {len(x)}")

# 使用Numpy的np.unique进行基准测试
start = perf_counter()
y1 = np.unique(x)
print(f"np.unique 耗时: {perf_counter() - start:.6f} 秒")

# 使用纯Python实现的位图法进行测试
start = perf_counter()
y2 = count_unique_bitmask(x)
print(f"位图法 (Python) 耗时: {perf_counter() - start:.6f} 秒")
print(f"结果一致性检查: {np.array_equal(y1, y2)}")

# ... (省略导入和RNG初始化)

@njit # 添加Numba JIT装饰器
def count_unique_bitmask_numba(ls):
    """
    使用Numba加速的位图法，对非负整数进行排序并去重。
    注意：此函数在特定条件下会遇到问题。
    """
    ret = []
    m = 0
    for x in ls:
        m = m | (1 << int(x))
    i = 0
    while m > 0:
        if (m & 1):
            ret.append(i)
        m = m >> 1
        i += 1
    return ret

# 再次测试Numba版本
start = perf_counter()
y3 = count_unique_bitmask_numba(x)
print(f"位图法 (Numba) 耗时: {perf_counter() - start:.6f} 秒")
# 此时np.array_equal(y1, y3) 将返回 False，因为y3是空列表
print(f"结果一致性检查 (Numba): {np.array_equal(y1, y3)}")

from numba import njit

@njit
def shift_test(amount):
    """
    演示在Numba中1左移操作对不同位数的行为，
    特别是在64位有符号整数中的符号位影响。
    """
    return 1 << amount

print("Numba中位移操作的行为:")
for i in range(66):
    result = shift_test(i)
    print(f"{i:2d}: {hex(result)} (十进制: {result})")
    if i == 63:
        print("  --- 注意: 1 << 63 在Numba int64中变为负数 ---")