如何在Numba njit环境下实现数组的itertools.product遍历计算？

因为itertools.product无法在Numba的@njit装饰函数中使用，下面提供两种实用的替代实现方案，适配不同场景：

方案1：嵌套循环实现（适合低维度场景） #

如果你的笛卡尔积维度较少（比如2-3个数组），直接用嵌套循环是最简单的方式，Numba对普通循环的优化效果很好。

示例代码：

import numba as nb
import numpy as np

@nb.njit
def product_2d(arr1, arr2):
    # 计算结果总长度
    total = len(arr1) * len(arr2)
    result = np.empty((total, 2), dtype=arr1.dtype)
    idx = 0
    for x in arr1:
        for y in arr2:
            result[idx] = (x, y)
            idx += 1
    return result

# 测试
a = np.array([1,2,3])
b = np.array([4,5])
print(product_2d(a, b))

如果是3个维度，只需再嵌套一层循环，逻辑完全一致。

方案2：通用笛卡尔积实现（支持任意维度） #

如果需要处理任意数量的输入数组，可以通过索引转换的方式实现通用笛卡尔积，核心是把线性索引转换成各维度的索引，再映射到原数组元素。

示例代码：

import numba as nb
import numpy as np

@nb.njit
def numba_product(arrays):
    # 获取每个数组的长度
    lengths = np.array([len(arr) for arr in arrays], dtype=np.int64)
    # 计算各维度的步长（类似进制转换的基数）
    strides = np.empty_like(lengths)
    strides[-1] = 1
    for i in range(len(lengths)-2, -1, -1):
        strides[i] = strides[i+1] * lengths[i+1]
    # 总元素数量
    total = strides[0] * lengths[0]
    # 初始化结果数组
    result = np.empty((total, len(arrays)), dtype=arrays[0].dtype)
    
    for i in range(total):
        current = i
        for j in range(len(lengths)):
            # 计算当前维度的索引
            dim_idx = current // strides[j]
            result[i, j] = arrays[j][dim_idx]
            current = current % strides[j]
    return result

# 测试
arrays = [np.array([1,2]), np.array([3,4,5]), np.array([6,7])]
print(numba_product(arrays))

注意事项 #

• 输入优先用numpy.ndarray类型，Numba对numpy数组的支持远优于Python原生列表
• 如果只需遍历组合无需存储结果，可以去掉结果数组的初始化，直接在循环内处理每个元素，节省内存
• 上述代码均可直接用@njit装饰，Numba会将其编译为机器码，性能接近原生C代码

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