Numba非Python模式下嵌套列表排序函数的替代实现需求

原代码中使用sorted配合lambda键函数的方式在Numba njit模式下不被支持，因为Numba对高阶函数和匿名函数的支持有限。以下是两种可行的替代实现：

方法一：基于Numpy数组的排序（推荐） #

利用Numba支持的np.lexsort实现多列排序，效率更高，适合大数据量场景：

import numba
import numpy as np

@numba.njit
def nested_sort_np(arr):
    # lexsort参数顺序：次排序键在前，主排序键在后
    # 此处实现先按第1列索引1升序，再按第2列索引2升序
    sort_indices = np.lexsort((arr[:, 2], arr[:, 1]))
    return arr[sort_indices]

# 测试示例
s = [[1, 3, 11], [2, 3, 19], [3, 2, 18], [4, 2, 9]]
# 将列表转换为numpy数组
input_arr = np.array(s)
# 调用排序函数并转回列表格式
sorted_result = nested_sort_np(input_arr).tolist()
print(sorted_result)
# 输出：[[4, 2, 9], [3, 2, 18], [1, 3, 11], [2, 3, 19]]

方法二：纯列表手动实现排序算法 #

如果不想依赖Numpy，可以手动实现排序逻辑（以冒泡排序为例），直接在循环中比较排序键：

import numba

@numba.njit
def nested_sort_list(s):
    # 复制输入列表避免修改原数据
    lst = [row.copy() for row in s]
    n = len(lst)
    for i in range(n):
        swapped = False
        for j in range(0, n - i - 1):
            # 先比较第1列，相等则比较第2列
            if lst[j][1] > lst[j+1][1] or (lst[j][1] == lst[j+1][1] and lst[j][2] > lst[j+1][2]):
                lst[j], lst[j+1] = lst[j+1], lst[j]
                swapped = True
        # 若某轮无交换，说明已排序完成，提前退出
        if not swapped:
            break
    return lst

# 测试示例
s = [[1, 3, 11], [2, 3, 19], [3, 2, 18], [4, 2, 9]]
print(nested_sort_list(s))
# 输出：[[4, 2, 9], [3, 2, 18], [1, 3, 11], [2, 3, 19]]

两种方案对比 #

• Numpy方案：运行效率高，Numba对Numpy操作的优化更充分，适合处理大规模数据。
• 纯列表方案：无需依赖Numpy，代码逻辑更贴近原生列表操作，适合小规模数据场景。

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