如何高效替换矩阵中除含特定值的行/列外所有值为np.nan？

我们需要实现的功能为：将输入矩阵X中除包含指定值v的行和列外，其余所有元素替换为np.nan，参考代码框架、输入示例与预期输出如下：

参考代码框架 #

import numpy as np
from numpy.typing import NDArray 


def get_masked_array(X: NDArray[float], v: float) -> NDArray[float]:
     # 待实现逻辑
     return arr
    

# 输入示例浮点数组
X = np.array([[ 1.,  2.,  2.,  3.,  3.],
              [ 1.,  2.,  2.,  4.,  4.],
              [ 5.,  5.,  6.,  6.,  6.],
              [ 7.,  8.,  9.,  9.,  9.],
              [10., 10., 10., 10., 10.]])

预期输出结果 #

>>> get_masked_array(X, 2.)
array([[ 1.,  2.,  2.,  3.,  3.],
       [ 1.,  2.,  2.,  4.,  4.],
       [nan,  5.,  6., nan, nan],
       [nan,  8.,  9., nan, nan],
       [nan, 10., 10., nan, nan]])

>>> get_masked_array(X, 3.)
array([[ 1.,  2.,  2.,  3.,  3.],
       [nan, nan, nan,  4.,  4.],
       [nan, nan, nan,  6.,  6.],
       [nan, nan, nan,  9.,  9.],
       [nan, nan, nan, 10., 10.]])

通过numpy布尔索引定位需要保留的行和列，其余位置赋值为np.nan即可，实现代码如下：

import numpy as np
from numpy.typing import NDArray 


def get_masked_array(X: NDArray[float], v: float) -> NDArray[float]:
    arr = X.copy()
    # 筛选所有包含指定值v的行索引
    keep_rows = (arr == v).any(axis=1)
    # 筛选所有包含指定值v的列索引
    keep_cols = (arr == v).any(axis=0)
    # 构建掩码：既不在保留行也不在保留列的位置需要设为nan
    mask = ~keep_rows[:, np.newaxis] & ~keep_cols
    arr[mask] = np.nan
    return arr

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