Abhishek Bhatia Abhishek Bhatia - 1 month ago 14
Python Question

Finding max in numpy skipping some rows and column

I want to find the max row and column index in a numpy matrix. But it not be in the a set of rows or columns. Thus, it should skip those rows and columns while computing the max.

Example:

# finding max in numpy matrix
[row,col] = np.where(mat == mat.max())


But it should skip rows removed_rows=[] and columns columns_rows=[]

I don't want to create a new sub matrix for the computation.

Answer

Let a be the input array, rows_rem and cols_rem be the rows and column indices to be skipped respectively. We would have an approach using masking, like so -

m,n = a.shape
d0,d1 = np.ogrid[:m,:n]
a_masked = a*~(np.in1d(d0,rows_rem)[:,None] | np.in1d(d1,cols_rem))
max_row, max_col = np.where(a_masked == a_masked.max())

Sample run -

In [204]: # Inputs
     ...: a = np.random.randint(11,99,(4,5))
     ...: rows_rem = [1,3]
     ...: cols_rem = [1,2,4]
     ...: 

In [205]: a
Out[205]: 
array([[36, 51, 72, 18, 31],
       [78, 42, 12, 71, 72],
       [38, 46, 42, 67, 12],
       [87, 56, 76, 14, 21]])

In [206]: a_masked
Out[206]: 
array([[64,  0,  0, 90,  0],
       [ 0,  0,  0,  0,  0],
       [17,  0,  0, 40,  0],
       [ 0,  0,  0,  0,  0]])

In [207]: max_row, max_col
Out[207]: (array([0]), array([3]))

Please note that if there's more than one element with the same max value, we would have all of those in the output. So, if you want any or the first of those, we can use argmax, like so -

max_row, max_col = np.unravel_index(a_masked.argmax(),a.shape)