如何根据条件删除Numpy数组axis 1元素并保留顺序？

核心思路 #

通过构造布尔掩码一次性标记所有需要保留的元素，避免拆分数组处理，严格维持原有的顺序，同时支持扩展到3D数组。

2D数组示例实现 #

以你提供的输入为例：

import numpy as np

# 输入数组
arr = np.array([[0,1,2,3],[0,2,3,4],[1,3,4,5]])

# 1. 确定每行需要移除的列索引：首元素为0时移除第2列，否则移除第3列
remove_col = np.where(arr[:, 0] == 0, 2, 3)

# 2. 创建全True的掩码，标记需要保留的元素
mask = np.ones_like(arr, dtype=bool)
# 将需要移除的位置设为False
mask[np.arange(len(arr)), remove_col] = False

# 3. 应用掩码并重塑为目标形状
result = arr[mask].reshape(len(arr), -1)
print(result)

输出：

[[0 1 3]
 [0 2 4]
 [1 3 4]]

3D数组扩展实现 #

假设3D数组形状为(深度, 行数, 列数)，处理逻辑和2D一致，仅需调整索引的广播方式：

import numpy as np

# 示例3D数组
arr_3d = np.array([
    [[0,1,2,3],[0,2,3,4],[1,3,4,5]],
    [[1,5,6,7],[0,8,9,10],[0,11,12,13]]
])

# 1. 确定每个行需要移除的列索引（形状：(深度, 行数)）
remove_col_3d = np.where(arr_3d[..., 0] == 0, 2, 3)

# 2. 创建全True的掩码（形状：(深度, 行数, 列数)）
mask_3d = np.ones_like(arr_3d, dtype=bool)
# 构造深度和行的索引，实现广播匹配
d_indices = np.arange(arr_3d.shape[0])[:, np.newaxis]
n_indices = np.arange(arr_3d.shape[1])
# 标记需要移除的位置
mask_3d[d_indices, n_indices, remove_col_3d] = False

# 3. 应用掩码并重塑为目标形状
result_3d = arr_3d[mask_3d].reshape(arr_3d.shape[0], arr_3d.shape[1], arr_3d.shape[2]-1)
print(result_3d)

输出：

[[[ 0  1  3]
  [ 0  2  4]
  [ 1  3  4]]

 [[ 1  5  6]
  [ 0  8 10]
  [ 0 11 13]]]

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