如何用向量化方式生成NumPy数组的多阶幂次数组？

算力阿峰

2026-5-28

当然有更高效的向量化方案！我们可以借助NumPy的广播机制或者累积乘积特性来替代显式循环，而且不同方法在不同场景下各有优势。先回顾下你的需求：

你有这样一个NumPy数组：

arr = np.array([[1, 2], [3, 4]])

想要生成一个包含该数组从0阶到指定order阶幂次的新数组，比如当order=3时，结果应该是：

arr_new = [[1, 1, 1, 2, 1, 4, 1, 8], [1, 1, 3, 4, 9, 16, 27, 64]]

你当前用for循环实现，但希望找到更快的向量化方法。

几种向量化实现方案

下面是三种高效的向量化方法，我们会结合基准测试来对比它们的性能：

1. 广播+Reshape 方案

利用NumPy的广播特性，直接让数组和幂次数组进行元素级运算，最后重塑成目标形状：

def broadcast_based_reshape(arr, order):
    # 创建幂次数组，通过广播匹配原数组维度
    powers = np.arange(order + 1)[:, None]
    # 计算各阶幂次，得到三维数组（行 × 幂次 × 列）
    exponentiated = arr[:, None] ** powers
    # 重塑为二维目标数组并返回
    return exponentiated.reshape(arr.shape[0], -1)

2. 广播+Hstack 方案

思路和上面类似，只是最后用np.hstack横向拼接各阶幂次的结果：

def broadcast_based_hstack(arr, order):
    # 创建三维幂次数组，触发广播机制
    powers = np.arange(order + 1)[:, None, None]
    # 计算各阶幂次
    exponentiated = arr ** powers
    # 横向拼接所有幂次结果
    return np.hstack(exponentiated)

3. 累积乘积(Cumprod) 方案

对于高阶幂次场景，累积乘积的方式可以避免重复计算幂次，性能提升更明显：

def broadcast_cumprod_reshape(arr, order):
    rows, cols = arr.shape
    # 预分配结果数组，中间维度存储0到order阶的幂次
    out = np.empty((rows, order + 1, cols), dtype=arr.dtype)
    # 0次幂恒为1
    out[:, 0, :] = 1
    # 将原数组广播到（行 × 阶数 × 列）的维度
    a = np.broadcast_to(arr[:, None], (rows, order, cols))
    # 通过累积乘积得到各阶幂次（从1阶开始）
    np.cumprod(a, axis=1, out=out[:,1:,:])
    # 重塑为二维目标数组
    return out.reshape(rows, -1)

基准测试对比

我们在两种典型场景下对比这些方法和你原来的循环实现：

小型数组：arr = np.array([[1, 2], [3,4]])，order=3
大型数组：arrLarge = np.random.randint(0, 10, (100, 100))，orderLarge=10

你的原循环实现代码：

def loop_based(arr, order):
    # 预分配结果数组
    rows, cols = arr.shape
    arr_new = np.zeros((rows, (order+1) * cols))
    # 遍历每个幂次，填充结果数组
    for i in range(order + 1):
        arr_new[:, (i * cols) : (i + 1) * cols] = arr**i
    return arr_new

测试结果

小型数组（2x2）：

%timeit -n 100000 loop_based(arr, order)
# 7.41 µs ± 174 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
%timeit -n 100000 broadcast_based_hstack(arr, order)
# 10.1 µs ± 137 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
%timeit -n 100000 broadcast_based_reshape(arr, order)
# 3.31 µs ± 61.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
%timeit -n 100000 broadcast_cumprod_reshape(arr, order)
# 11 µs ± 102 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

大型数组（100x100）：

%timeit -n 1000 loop_based(arrLarge, orderLarge)
# 261 µs ± 5.82 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit -n 1000 broadcast_based_hstack(arrLarge, orderLarge)
# 225 µs ± 4.15 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit -n 1000 broadcast_based_reshape(arrLarge, orderLarge)
# 223 µs ± 2.16 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit -n 1000 broadcast_cumprod_reshape(arrLarge, orderLarge)
# 157 µs ± 1.02 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

结论

处理小型数组时，广播+Reshape的方法速度最快，比原循环实现快了一倍多；
处理大型数组时，累积乘积(Cumprod)的方法扩展性最好，性能领先其他方案，因为它避免了重复计算高次幂，效率更高。

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