使用 Numba 中断在线性时间内对非负整数进行排序

Question

我正在尝试在线性时间内对非负整数数组/列表进行排序。我们也只保留独特的元素。这是一个例子，

Sort: [7, 7, 0, 3, 2, 1, 9, 1]

7: 10000000
7: 10000000
0: 10000001
3: 10001001
2: 10001101
1: 10001111
9: 1010001111
1: 1010001111

1010001111: []
101000111: [0]
10100011: [0, 1]
1010001: [0, 1, 2]
101000: [0, 1, 2, 3]
10100: [0, 1, 2, 3]
1010: [0, 1, 2, 3]
101: [0, 1, 2, 3]
10: [0, 1, 2, 3, 7]
1: [0, 1, 2, 3, 7]
: [0, 1, 2, 3, 7, 9]

本质上，我是在线性时间内实现

np.unique([7, 7, 0, 3, 2, 1, 9, 1])

。这是我的Python，

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

# @njit
def count(ls):
    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

RNG = np.random.default_rng(0)

x = RNG.integers(2**16, size=2**17)
start = perf_counter()
y1 = np.unique(x)
print(perf_counter() - start)
start = perf_counter()
y2 = count(x)
print(perf_counter() - start)
print((y1 == y2).all())

我的“O(n)”排序没有击败 Numpy 的独特功能。我预计，因为 Python 比 C 慢（我猜这就是实现

np.unique

的地方）。为了解决这个问题，我尝试使用 Numba 的 JIT 装饰器。但是，如果我取消注释装饰器，函数会以某种方式中断并返回一个空列表。它无需装饰器即可工作。

有人可以指出我的疏忽吗？

Answer 1

Python 使用任意精度整数来表示所有整数。这通常很有用，但缺点是速度慢。为了加快速度，Numba 使用 64 位有符号整数。

这样做的后果之一是 1 << 63 is a negative number in Numba.

以下测试程序将显示这一点。

from numba import njit


@njit
def shift(amount):
    return 1 << amount


for i in range(66):
    print(i, hex(shift(i)))

然后，

while m > 0:

立即退出。这就是为什么 numba 版本给出一个空列表，也是为什么如果您有超过 64 个数字，它通常无法正常工作。