2D numpy掩码未按预期工作：将2x3数组转为2x2数组遇阻

嘿，我懂你碰到的这个问题——想用掩码数组给2x3的numpy数组每行删一个元素变成2x2，但结果总是不如预期，甚至变成了一维数组？这其实是因为numpy的掩码索引规则和你想的不太一样，咱们一步步来解决它。

先搞清楚为什么掩码方法会“翻车” #

假设你有这样一个2x3数组：

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

你可能尝试过这样写掩码：

# 第一行保留索引0和2，第二行保留0和1
mask = np.array([[True, False, True], [True, True, False]])
result = arr[mask]
print(result)  # 输出：array([1, 3, 4, 5])

结果是个一维数组，而不是你想要的2x2。这是因为numpy的掩码索引会把所有标记为True的元素提取出来，自动扁平化结果，不会保留原来的行结构——它只关心哪些元素符合条件，不关心它们属于哪一行。

正确的做法：用索引数组精准提取 #

要实现“每行移除指定索引，保持2x2结构”，更直接的方式是构造每行要保留的索引数组，然后结合行索引的广播来提取元素。

步骤1：定义每行要移除的索引 #

比如你想第一行删索引1，第二行删索引2：

remove_indices = [1, 2]  # 对应每行要移除的列索引

步骤2：生成每行要保留的索引 #

我们可以通过集合操作或者列表推导来生成每行要保留的列索引：

all_cols = np.arange(arr.shape[1])  # 所有列索引：[0,1,2]
# 对每个要移除的索引，生成剩下的列索引
keep_indices = np.array([np.setdiff1d(all_cols, [idx], assume_unique=True) for idx in remove_indices])
# 此时keep_indices是：array([[0,2], [0,1]])

步骤3：广播行索引，提取目标元素 #

用np.arange(arr.shape[0])[:, None]把行索引转换成列向量，这样就能和keep_indices广播匹配，每行对应自己的保留索引：

result = arr[np.arange(arr.shape[0])[:, None], keep_indices]
print(result)
# 输出正是你想要的2x2数组：
# [[1 3]
#  [4 5]]

更简洁的场景处理 #

如果你的需求是每行移除固定位置的元素（比如都删最后一列），那直接切片就搞定了：

result = arr[:, :-1]  # 移除每行最后一个元素，得到2x2数组

但如果是每行移除不同的索引，上面的索引数组方法就是最通用的解决方案啦。

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