Modify every two largest elements of matrix rows and columns

Question

In python, I have a matrix and I want to find the two largest elements in every row and every column and change their values to 1 (seperately, I mean get two matrices where one of them modified the rows and the other modified the cols). The main goal is to get a corresponding matrix with zeros everywhere except those ones I've put in the 2 largest element of each row and column (using np.where(mat == 1, 1, 0). I'm trying to use the np.argpartition in order to do so but without success. Please help. See image below.

Answer 1

Here's an approach with np.argpartition -

idx_row = np.argpartition(-a,2,axis=1)[:,:2]
out_row = np.zeros(a.shape,dtype=int)
out_row[np.arange(idx_row.shape[0])[:,None],idx_row] = 1

idx_col = np.argpartition(-a,2,axis=0)[:2]
out_col = np.zeros(a.shape,dtype=int)
out_col[idx_col,np.arange(idx_col.shape[1])] = 1

Sample input, output -

In [40]: a
Out[40]: 
array([[ 3,  7,  1, -5, 14,  2,  8],
       [ 5,  8,  1,  4, -3,  3, 10],
       [11,  3,  5,  1,  9,  2,  5],
       [ 6,  4, 12,  6,  1, 15,  4],
       [ 8,  2,  0,  1, -2,  3,  5]])

In [41]: out_row
Out[41]: 
array([[0, 0, 0, 0, 1, 0, 1],
       [0, 1, 0, 0, 0, 0, 1],
       [1, 0, 0, 0, 1, 0, 0],
       [0, 0, 1, 0, 0, 1, 0],
       [1, 0, 0, 0, 0, 0, 1]])

In [42]: out_col
Out[42]: 
array([[0, 1, 0, 0, 1, 0, 1],
       [0, 1, 0, 1, 0, 1, 1],
       [1, 0, 1, 0, 1, 0, 0],
       [0, 0, 1, 1, 0, 1, 0],
       [1, 0, 0, 0, 0, 0, 0]])

---

Alternatively, if you are into compact codes, we can skip the initialization and use broadcasting to get the outputs from idx_row and idx_col directly, like so -

out_row = (idx_row[...,None] == np.arange(a.shape[1])).any(1).astype(int)
out_col = (idx_col[...,None] == np.arange(a.shape[0])).any(0).astype(int).T