高效移除行排列整数矩阵中的指定整数k

现有一个n×n矩阵，每行都是整数1到n的排列。给定整数k∈{1,2,...,n}，需要定位每行中的k并移除，同时保持每行剩余元素的原有顺序。

以n=5的矩阵X为例：

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

当k=1时，期望输出为：

X = np.array([
    [4, 3, 5, 2],
    [5, 3, 2, 4],
    [3, 2, 4, 5],
    [2, 4, 3, 5],
    [3, 5, 2, 4]
])

当前采用逐行循环的实现方式，小矩阵场景可行，但处理大矩阵时扩展性差：

m = 5
n = m
X = np.array([np.random.permutation(np.arange(1, n + 1)) for _ in range(m)])
k = 1
mask = X != k
Xnew = np.zeros([n,m-1])
for i in range(m):
    Xnew[i]=X[i][mask[i]]

利用numpy的矢量化操作替代循环，大幅提升大矩阵场景下的处理效率，以下是几种可行方案：

方法一：直接通过布尔掩码索引重塑 #

由于每行恰好包含一个k，掩码X != k每行有且仅有n-1个True值，直接索引后重塑形状即可：

import numpy as np

n = 5
X = np.array([
    [4, 3, 5, 1, 2],
    [5, 3, 2, 1, 4],
    [3, 2, 4, 1, 5],
    [2, 4, 3, 5, 1],
    [3, 5, 2, 1, 4]
])
k = 1

mask = X != k
Xnew = X[mask].reshape(n, n-1)
print(Xnew)

方法二：使用np.compress按行压缩 #

调用np.compress指定按行维度压缩，保留掩码为True的元素：

mask = X != k
Xnew = np.compress(mask, X, axis=1).reshape(n, n-1)

方法三：定位k的位置后删除 #

先找到每行中k的索引，再构造排除该位置的索引数组取值：

# 获取每行k所在的列索引
k_col_indices = np.argwhere(X == k)[:, 1]
# 构造行索引：每行重复n-1次
row_indices = np.repeat(np.arange(n), n-1)
# 构造列索引：每行排除k的位置后扁平化
col_indices = np.array([np.delete(np.arange(n), idx) for idx in k_col_indices]).flatten()
# 索引取值并重塑形状
Xnew = X[row_indices, col_indices].reshape(n, n-1)

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