如何在NumPy中为每行使用不同切片进行二维索引？

你遇到的问题很典型——NumPy的普通切片语法没法直接给每行指定不同的切片范围，所以才会抛出那个类型错误。下面给你几种可行的解决方案，顺便聊聊各自的性能特点：

方法一：花式索引（灵活适配任意列选择） #

如果你的列选择不一定是连续的，或者需要更灵活的索引规则，可以直接构造每个目标元素的行和列索引，用NumPy的花式索引取出：

import numpy as np

x = np.arange(12).reshape(3, 4)
# 构造行索引：前两行，每行对应2个元素
row_idx = np.repeat([0, 1], 2)
# 构造列索引：第一行取0、1，第二行取1、2
col_idx = np.array([0, 1, 1, 2])
# 取出元素后重新reshape成目标形状
y = x[row_idx, col_idx].reshape(2, 2)
print(y)
# 输出：
# [[0 1]
#  [5 6]]

这种方法的好处是完全灵活，不管每行要取哪些列（连续或不连续）都能搞定，唯一的小缺点是需要手动构造索引数组，行数多的时候要注意索引的正确性。

方法二：用np.take_along_axis（适合连续切片场景） #

如果你的需求是每行取连续的N列，只是起始位置不同，那take_along_axis会是最简洁高效的选择：

import numpy as np

x = np.arange(12).reshape(3, 4)
# 定义每行切片的起始列索引
starts = np.array([0, 1])
# 生成每行要取的列索引：每行取2个连续列
cols = starts[:, None] + np.arange(2)
# 直接提取对应位置的元素
y = np.take_along_axis(x, cols, axis=1)
print(y)
# 输出：
# [[0 1]
#  [5 6]]

这个方法是NumPy原生优化的矢量操作，不需要手动处理行索引，代码更简洁，性能也最优，推荐在连续切片的场景下优先使用。

方法三：列表推导式（小数据量友好） #

如果你的数据集不大，追求代码的直观性，用列表推导式循环处理每行也是可以的：

import numpy as np

x = np.arange(12).reshape(3, 4)
# 给每行指定对应的切片
slices = [slice(0, 2), slice(1, 3)]
# 循环取出每行的切片，再合并成二维数组
y = np.vstack([x[i, s] for i, s in enumerate(slices)])
print(y)
# 输出：
# [[0 1]
#  [5 6]]

这种方法读起来一目了然，但本质是Python层面的循环，当处理超大数组（比如上万行以上）时，性能会比前两种NumPy原生方法差很多，适合小数据量的场景。

性能对比 #

• 列表推导式：最慢，Python循环的开销会在大数据量下被放大，不推荐用于大规模计算。
• 花式索引：性能不错，但需要额外存储两个索引数组，内存占用略高，不过绝大多数场景下可以忽略。
• take_along_axis：性能最优，内存效率也高，是连续切片场景下的首选。

备注：内容来源于stack exchange，提问作者galah92