如何在NumPy中利用起始/结束索引数组对二维数组逐行切片？

给定4行6列的NumPy数组，以及每行对应的起始索引数组starts和结束索引数组ends，要从每行中提取[start, end)区间的3个元素，常规切片array[starts:ends]无法实现（该语法是对行维度切片，而非逐行切列），以下是两种可行方案：

方案一：广播生成索引+花式索引 #

先生成每行对应的列索引，结合行索引完成花式索引：

import numpy as np

# 原始数组
array = np.array([[3, 5, 2, 5, 1, 3],
                  [5, 3, 0, 1, 4, 3],
                  [2, 1, 0, 2, 2, 4],
                  [2, 1, 0, 4, 2, 2]])
starts = np.array([3, 0, 1, 0])  # 由示例ends反推：ends = starts + 3
ends = np.array([6, 3, 4, 3])

# 生成每行对应的3个列索引
col_indices = starts[:, None] + np.arange(3)
# 生成匹配的行索引
row_indices = np.arange(array.shape[0])[:, None]

# 执行花式索引
result = array[row_indices, col_indices]
print(result)

输出：

array([[5, 1, 3],
       [5, 3, 0],
       [1, 0, 2],
       [2, 1, 0]])

方案二：使用np.take_along_axis（更简洁） #

利用NumPy专用函数take_along_axis，沿列方向按指定索引提取元素：

import numpy as np

array = np.array([[3, 5, 2, 5, 1, 3],
                  [5, 3, 0, 1, 4, 3],
                  [2, 1, 0, 2, 2, 4],
                  [2, 1, 0, 4, 2, 2]])
starts = np.array([3, 0, 1, 0])
ends = np.array([6, 3, 4, 3])

col_indices = starts[:, None] + np.arange(3)
# 沿轴1（列维度）提取对应索引的元素
result = np.take_along_axis(array, col_indices, axis=1)
print(result)

输出与方案一完全一致。

核心说明 #

常规切片array[starts:ends]是对行维度进行范围切片，无法实现"每行按不同列区间截取"的需求。上述两种方案的核心是为每行生成专属的列索引集合，再通过索引匹配完成逐行的列切片操作。

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