如何优化二维NumPy数组逐行应用带滚动n参数的lead函数？

你的需求是对(121,121)的二维数组Qxx的每一行，生成对应n从0到120的变换结果，最终得到(121,121,121)的三维数组。嵌套列表推导式慢的核心原因是Python循环的开销，推荐用NumPy向量化操作替代，以下是两种高效方案：

方案1：简化逻辑+向量化循环 #

首先可以简化你的lead函数，因为np.roll后替换最后n位的操作，等价于直接取数组的第n位到末尾，再拼接n个1：

def lead(x, n):
    if n == 0:
        return x.copy()
    return np.concatenate([x[n:], np.ones(n, dtype=x.dtype)])

然后用NumPy的向量化赋值替代嵌套循环，只需要循环n的取值，对所有行批量处理：

import numpy as np

L = Qxx.shape[0]  # 121
result = np.ones((L, L, L), dtype=Qxx.dtype)

for n in range(L):
    # 对所有行，将第n组结果的前(L-n)列替换为Qxx对应行的[n:]切片
    result[:, n, :L - n] = Qxx[:, n:]

这个方案比嵌套列表推导式快很多，因为循环仅针对n，且每行的赋值是NumPy内部的向量化操作，避免了Python层面的逐元素循环。

方案2：完全向量化（无Python循环） #

利用NumPy的广播和高级索引，实现完全无Python循环的操作，效率最高：

import numpy as np

L = Qxx.shape[0]
# 生成索引矩阵：shape=(1, L, L)，每个位置对应n + k（n是当前变换的参数，k是结果数组的列索引）
indices = np.arange(L)[None, :, None] + np.arange(L)[None, None, :]
# 生成掩码：标记哪些位置需要填充为1（即n + k >= L的位置）
mask = indices >= L

# 利用take获取对应位置的原数组值，再用where替换掩码位置为1
result = np.where(
    mask,
    1,
    Qxx[:, None, :].take(indices % L, axis=2)
)

逻辑说明： #

• indices通过广播生成所有n和k的和，indices % L对应np.roll后的索引；
• mask标记需要设为1的位置（即原roll后的最后n位）；
• Qxx[:, None, :].take(...)将原数组扩展为三维后，按索引批量取值；
• 最后用np.where完成填充1的操作。

关于apply_along_axis的说明 #

你尝试的apply_along_axis写法无法直接实现需求：因为apply_along_axis是对每一行单独处理一次，每次只能传递一个n参数，无法一次性为每行生成所有n对应的结果。它的底层本质还是Python循环，效率提升有限，因此不推荐使用。

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