如何在NumPy多维数组的每行中移除指定索引的元素？

给定以下NumPy数组：

import numpy as np

seed_idx = np.asarray([[0, 1],
                       [1, 2],
                       [2, 3],
                       [3, 4]], dtype=np.int_)

target_idx = np.asarray([[2,9,4,1,8],
                         [9,7,6,2,4],
                         [1,0,0,4,9],
                         [7,1,2,3,8]], dtype=np.int_)

需求：对target_idx的每一行，选取不在对应行seed_idx中的元素，预期结果为：

[[4,1,8],
 [9,2,4],
 [1,0,9],
 [7,1,2]]

即实现np.take_along_axis(target_idx, seed_idx, axis=1)的反向操作——排除指定索引元素而非保留。

优雅实现方案 #

方法1：逐行掩码筛选（高效适配大数据量） #

通过生成布尔掩码标记需要保留的索引，适合大规模数组场景：

# 生成每行的完整索引序列
row_indices = np.arange(target_idx.shape[1])[np.newaxis, :]
# 生成掩码：标记索引不在seed_idx对应行中的位置
mask = ~np.isin(row_indices, seed_idx)
# 应用掩码筛选并重塑结果形状
result = target_idx[mask].reshape(target_idx.shape[0], -1)

方法2：列表推导结合np.delete（代码极简） #

代码更简洁直观，适合中小规模数组：

result = np.array([np.delete(row, idx) for row, idx in zip(target_idx, seed_idx)])

两种方法均可得到符合预期的结果，可根据数组规模选择合适方案。

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