如何优雅地对NumPy数组执行逐行滚动操作？

默认的np.roll无法实现每行使用不同偏移量的滚动操作，这里提供两种高效的向量化解决方案：

方法一：构造偏移索引实现 #

通过为每行生成对应的列索引偏移，利用NumPy的索引机制直接提取元素，完全向量化，效率拉满。

import numpy as np

arr = np.array([[1, 2, 3, 4],
                [5, 6, 7, 8]])
shifts = np.array([1, 2])

n_rows, n_cols = arr.shape
# 生成基础列索引
cols = np.arange(n_cols)
# 计算每行偏移后的列索引（右移shift等价于索引减shift后取模）
cols_shifted = (cols - shifts[:, np.newaxis]) % n_cols
# 通过索引获取结果
result = arr[np.arange(n_rows)[:, np.newaxis], cols_shifted]

print(result)
# 输出：
# [[4 1 2 3]
#  [7 8 5 6]]

方法二：使用np.take简化代码 #

用np.take替代手动索引，代码更简洁，效果完全一致：

result = np.take(arr, cols_shifted, axis=1)

这两种方法都避免了显式循环，充分利用NumPy的向量化特性，处理大数组时性能远优于循环实现。

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