如何用纯NumPy实现从二维数组每行提取可变数量前N元素并拼接为一维数组？

当然有纯NumPy的实现方案，而且代码简洁高效，无需循环操作。

实现步骤与代码 #

先定义示例中的输入数组：

import numpy as np

# 二维输入数组
arr_2d = np.array([[22, 66, 16, 17, 79],
                   [17,  5, 27, 58, 79],
                   [ 7, 41, 66, 82, 48]])
# 指定每行保留的元素数量
keep_counts = np.array([3, 4, 1])

通过广播生成布尔掩码的方式提取目标元素：

# 生成列索引数组，形状为(5,)
col_indices = np.arange(arr_2d.shape[1])
# 利用广播让每行的保留数与列索引比较，生成形状为(3,5)的布尔掩码
mask = col_indices < keep_counts[:, np.newaxis]
# 通过掩码提取元素，自动拼接为一维数组
result = arr_2d[mask]

执行后result的结果为：

array([22, 66, 16, 17,  5, 27, 58,  7])

原理说明 #

1. np.arange(arr_2d.shape[1])生成从0到列数-1的索引，对应二维数组的每一列位置；
2. keep_counts[:, np.newaxis]把一维的保留数数组转换为列向量（形状从(3,)变为(3,1)），借助NumPy的广播特性，让它和列索引数组逐行比较；
3. 生成的布尔掩码mask中，True位置对应需要保留的元素，直接用该掩码索引二维数组，即可提取所有目标元素并自动展平为一维数组。

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