NumPy二维数组如何对不同行执行不等长度的切片操作

首先明确核心前提：NumPy常规数值型二维数组要求所有行长度完全一致，这是由其底层连续内存的存储设计决定的，因此不存在单次切片语法直接返回规则二维不等长数组的可能，所有不等长切片的结果要么是Python列表，要么是object类型的NumPy数组。

基于示例数组：

import numpy as np
a = np.array([[1,2,3],
              [4,5,6]])

需求为第一行取全部3个元素、第二行取前2个元素，可通过以下两种主流方式实现：

方法1：逐行切片组装（日常场景首选） #

写法直观无理解成本，代码可维护性高，绝大多数场景下性能够用：

# 按顺序定义每一行从行首开始取的元素个数
take_counts = [3, 2]
res = [row[:cnt] for row, cnt in zip(a, take_counts)]

运行后res是存储各行列切片结果的列表，内容为[array([1, 2, 3]), array([4, 5])]。
如果需要自定义更灵活的切片规则（比如不是从行首开始取），可以直接传入每个行对应的slice对象：

# 示例：第一行取全部、第二行取索引1到末尾的元素
slices = [slice(None), slice(1, None)]
res = [row[s] for row, s in zip(a, slices)]

如果需要把结果转为NumPy数组，指定dtype=object即可：

res_arr = np.array(res, dtype=object)

方法2：向量化掩码实现（大数组性能优化场景） #

如果处理的数组规模极大（十万行以上级别），Python层面的逐行循环会有性能损耗，可以用NumPy向量化操作实现，速度会快数倍到数十倍：

take_counts = [3, 2]
# 生成列索引矩阵，对比每行要取的元素个数生成保留掩码
col_indices = np.arange(a.shape[1])
mask = col_indices < np.array(take_counts)[:, None]
# 取出所有符合条件的元素后，按每行的元素个数拆分
res = np.split(a[mask], np.cumsum(take_counts)[:-1])

注意：不要为了炫技强行用复杂技巧，数据量不大的时候优先选方法1，代码出问题更容易排查。

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