如何高效提取二维Numpy数组每行对应指定数量的前N个元素？

问题场景 #

现有一个M×N的Numpy数组，需要为每行提取前x个元素，其中x是长度为M的数组（或列表）。例如：

import numpy as np
M = 5
N = 4
a = np.ones((M, N))
x = np.array([2, 3, 1, 4, 2])  # 每行要提取的元素个数

期望输出为一维数组组成的列表（避免Numpy不推荐的不规则二维数组）：

[array([1., 1.]),
 array([1., 1., 1.]),
 array([1.]),
 array([1., 1., 1., 1.]),
 array([1., 1.])]

直接使用a[:, :x]会报错，因为索引不支持数组输入，同时需要避免低效的显式循环。

解决方案1：列表推导式（简洁高效） #

利用Python优化后的列表推导式实现，代码简洁且性能优于手动循环：

result = [a[i, :x[i]] for i in range(M)]

• 优势：代码直观易懂，Python内部对列表推导式做了高度优化，对于绝大多数大数据场景性能足够。
• 输出：直接得到符合要求的一维数组列表，无需额外处理。

解决方案2：向量化掩码拆分（超大数组优化） #

如果数组规模极大，可采用全向量化操作先提取所有目标元素，再拆分得到结果：

# 生成列索引矩阵
col_indices = np.tile(np.arange(N), (M, 1))
# 生成掩码：标记每行中需要保留的元素
mask = col_indices < x[:, np.newaxis]
# 提取所有符合条件的元素，再按x的长度拆分
flat_elements = a[mask]
result = np.split(flat_elements, np.cumsum(x)[:-1])

• 原理：先通过掩码筛选出所有需要保留的元素，再利用np.cumsum计算拆分位置，最后用np.split拆分为对应长度的数组列表。
• 优势：全程使用Numpy向量化操作，避免Python层面的循环，在超大规模数组下性能更优。

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