[FIXED] What is an efficient way to replace all values of a matrix except for those in rows or columns which contain a specific value?

Issue

I'm trying to replace all values of the input matrix X with np.nan except for the rows and columns which contain a value v:

import numpy as np
from numpy.typing import NDArray 


def get_masked_array(X: NDArray[float], v: float) -> NDArray[float]:
     # something
     return arr
    

# Input float array
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.]])

Expected results:

>>> 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.]])

Solution

Check row and column condition first and then combine them into a boolean condition taking advantage of numpy broadcasting:

v = 2

eq_v = X == v
has_v = eq_v.any(0) | eq_v.any(1, keepdims=True)  # check if row or col has v
np.where(has_v, X, np.nan)

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]])

Using walrus operator (>= python 3.8), you can do this in one line as efficiently:

np.where((eq_v := (X == v)).any(0) | eq_v.any(1, keepdims=True), X, np.nan)
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]])

Answered By - Psidom