如何用NumPy实现二维数组按一维数组移位？三维数组按二维数组移位同理

问题描述 #

给定二维数组a和一维数组b：

import numpy as np
a = np.array([[1,1,1,1],[2,2,2,2],[3,3,3,3]])
b = np.array([0,2,1,0])

需要得到如下结果：

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

核心需求是让数组a的元素沿指定轴，按b的对应值进行循环移位，且需处理大型三维数组（约(10000,10000,100)），必须避免Python循环以保证效率。

解决方案：矢量化索引实现 #

利用NumPy的高级索引和广播机制，可以完全实现无循环的高效移位操作，核心思路是构造移位后的索引矩阵，直接通过索引提取目标元素。

二维数组实现代码 #

import numpy as np

a = np.array([[1,1,1,1],[2,2,2,2],[3,3,3,3]])
b = np.array([0,2,1,0])

# 构造行索引网格，适配列维度的广播
rows = np.arange(a.shape[0])[:, np.newaxis]
# 计算循环移位后的行索引：(原行索引 - 对应列移位量) 对行数取模
shifted_rows = (rows - b) % a.shape[0]
# 高级索引提取结果，列索引保持原顺序
out = a[shifted_rows, np.arange(a.shape[1])]

运行后out完全符合预期结果。

三维数组扩展 #

针对形状为(N, M, K)的三维数组，若需沿轴0（第一个维度），按长度为K的数组b对每个第三维元素进行移位，代码如下：

import numpy as np

# 示例三维数组
a = np.random.rand(10000, 10000, 100)
b = np.random.randint(0, 10000, size=100)

# 构造三维行索引网格
rows = np.arange(a.shape[0])[:, np.newaxis, np.newaxis]
# 计算移位后的行索引
shifted_rows = (rows - b) % a.shape[0]
# 构造列和第三维的索引网格，适配广播
cols = np.arange(a.shape[1])[np.newaxis, :, np.newaxis]
depth_idx = np.arange(a.shape[2])[np.newaxis, np.newaxis, :]
# 提取结果
out = a[shifted_rows, cols, depth_idx]

该方法完全基于NumPy的矢量化操作，底层由C实现，处理大型数组时效率远高于Python循环。

移位方向调整 #

如果需要向上循环移位而非向下，只需将索引计算改为：

shifted_rows = (rows + b) % a.shape[0]

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