如何用索引快速选取二维NumPy数组子数组的非第三个元素？

给定如下二维NumPy数组：

import numpy as np

array = np.array([[1, 1, 4, 2, 1, 8],
                  [1, 1, 8, 2, 1, 16],
                  [1, 1, 40, 2, 1, 80],
                  [1, 2, 40, 2, 1, 80]])

需要移除每个子数组中的每第三个元素，得到结果：

[[1 1 2 1]
 [1 1 2 1]
 [1 1 2 1]
 [1 2 2 1]]

以下是两种无需循环和np.concatenate的高效实现方法，均基于NumPy原生向量化操作，适合处理大维度数组：

方法一：布尔掩码索引 #

通过生成重复的布尔掩码，直接筛选保留的元素：

# 生成掩码：每3个元素保留前两个，重复至子数组长度
mask = np.tile([True, True, False], array.shape[1] // 3)
result = array[:, mask]

方法二：Reshape + 切片（更高效） #

利用数组形状变换，直接提取每个3元素块的前两个元素：

# 将每行重塑为(块数, 3)，取前两列后再展平回原行数
result = array.reshape(array.shape[0], -1, 3)[:, :, :2].reshape(array.shape[0], -1)

效率说明 #

两种方法均为向量化操作，无Python级循环：

• 方法二的reshape在数组内存连续时为视图操作（仅修改数组元数据，不复制内存），切片也直接操作底层数据，整体性能更优，尤其适合超大第一维度的数组。
• 方法一的掩码生成和索引也足够高效，代码可读性更强。

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