如何无循环实现Numpy二维数组按对应行偏移量滚动？

给定二维数组和对应行的偏移量：

import numpy as np
A = np.array([[0,0,1,1,0,0,0],[0,1,1,0,1,0,0],[1,0,0,1,1,0,0]])
shifts = np.array([2,4,1])

要实现每行按对应偏移量滚动且不使用for循环，可通过向量化索引高效完成：

n_rows, n_cols = A.shape
# 生成基础列索引
base_cols = np.arange(n_cols)
# 计算每行的目标列索引（广播运算+取模保证合法）
target_cols = (base_cols - shifts[:, np.newaxis]) % n_cols
# 索引得到滚动结果
rolled_A = A[np.arange(n_rows)[:, np.newaxis], target_cols]

输出结果：

print(rolled_A)
# [[0 0 0 0 0 1 1]
#  [1 0 0 0 1 1 0]
#  [0 0 1 1 0 0 1]]

核心逻辑 #

• shifts[:, np.newaxis]将一维偏移量转为列向量，实现和基础列索引的广播运算
• 用取模% n_cols确保偏移后的索引始终在数组列的合法范围内
• 通过行索引矩阵np.arange(n_rows)[:, np.newaxis]和目标列索引矩阵组合，直接完成所有行的批量索引，完全规避Python循环

这种方案依赖Numpy底层的C实现，处理大规模数组时性能远优于循环写法。

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