Python中无需循环从多维数组按指定起始点提取固定长度子数组的方法

你遇到的报错本质是Python原生切片语法不支持传入数组作为起止参数，仅支持标量输入，因此array1[:, :, array2:array2 + 50]的写法无法生效。

最优实现方案 #

用np.take_along_axis实现，代码全程无Python级循环，是完全矢量化的高效实现：

# 生成每个位置对应要提取的50个索引，shape自动变为(M, N, 50)
extract_indices = array2 + np.arange(50)
# 沿最后一个轴按生成的索引提取数据
result = np.take_along_axis(array1, extract_indices, axis=2)

可选兼容方案（适配无take_along_axis的旧版本NumPy） #

通过广播构造三维高级索引实现相同效果：

# 构造前两维的广播索引
m_axis_idx, n_axis_idx = np.indices(array1.shape[:2], sparse=True)
extract_indices = array2 + np.arange(50)
result = array1[m_axis_idx, n_axis_idx, extract_indices]

注意事项 #

• 提取前建议先对array2做范围截断，避免索引越界：

  max_valid_start = array1.shape[2] - 50
  array2 = np.clip(array2, 0, max_valid_start)
  

• 两种方案输出的shape均为预期的(M, N, 50)，在你给出的M=10、N=5、P=2000的场景下，运行效率远高于循环遍历实现。

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