如何改写Numba njit函数以支持任意数量数组的组合求和？

问题根源 #

1. 递归过程中未正确处理回溯逻辑，导致索引组合混乱，输出不符合预期
2. Numba List初始化采用临时方案（如未明确指定元素类型），引发类型不兼容或数据丢失

修复后的完整代码 #

import numba
from numba import njit
from numba.typed import List

@njit
def _recursive_search(arrays, target, current_indices, current_sum, result):
    # 终止条件：遍历完所有数组
    if len(current_indices) == len(arrays):
        if current_sum == target:
            # 拷贝当前索引组合到结果列表
            result.append(numba.typed.List(current_indices))
        return
    
    current_arr = arrays[len(current_indices)]
    for idx in range(len(current_arr)):
        new_sum = current_sum + current_arr[idx]
        # 剪枝优化（仅适用于非负元素数组）
        if new_sum > target:
            continue
        # 递归前添加当前索引
        current_indices.append(idx)
        _recursive_search(arrays, target, current_indices, new_sum, result)
        # 回溯：移除当前索引，恢复状态
        current_indices.pop()

@njit
def find_combinations(arrays, target):
    # 正确初始化Numba List：存储int类型索引的列表
    result = List()
    current_indices = List.empty_list(numba.int64)
    _recursive_search(arrays, target, current_indices, 0, result)
    return result

关键修复说明 #

• 递归回溯修正：每次递归返回后弹出当前索引，确保下一轮遍历的索引列表状态正确，避免索引组合错乱
• Numba List初始化修复：使用List.empty_list(numba.int64)明确指定索引列表的元素类型，替代临时方案，解决类型推导错误
• 剪枝优化：添加提前终止逻辑，当累加和超过目标值时跳过后续遍历，提升性能（若数组含负数可移除该判断）

使用示例 #

# 构造任意长度的Numba typed数组列表
test_arrays = List()
test_arrays.append(numba.typed.List([1, 2, 3]))
test_arrays.append(numba.typed.List([4, 5]))
test_arrays.append(numba.typed.List([6, 7]))

# 查找元素和为12的索引组合
target_sum = 12
matches = find_combinations(test_arrays, target_sum)

# 打印结果
for combo in matches:
    print([int(i) for i in combo])
# 输出：
# [0, 1, 0]  # 1 + 5 + 6 = 12
# [1, 0, 0]  # 2 + 4 + 6 = 12

注意事项 #

• 所有输入数组必须转换为Numba typed List，原生Python列表无法被Numba njit编译处理
• 若数组长度极大（超过Numba默认递归深度限制），建议改用迭代实现，避免栈溢出
• 若数组包含负数，需移除剪枝逻辑，否则会漏掉有效组合

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