Matrix row masking at different indices

I have an 2-D array that I want to find the max value per row and then find the next max-value that is not within +/- n of the previous value. For example I have the following matrix:

results = 
array([[ 33, 108, 208,  96,  96, 112,  18, 208,  33, 323,  60,  42],
       [ 51,   6,  39, 112, 160, 144, 342, 195,  27, 136,  42,  54],
       [ 12, 176, 266, 162,  45,  70, 156, 198, 143,  56, 342, 130],
       [ 22, 288, 304, 162,  21, 238, 156, 126, 165,  91, 144, 130],
       [342, 120,  36,  51,  10, 128, 156, 272,  32,  98, 192, 288]])

row_max_index = results.argmax(1)
row_max_index #show max index
array([ 9,  6, 10,  2,  0])

Now I’d like to get the next max value not within say +/- 2 of the current max.

Here is what I have but it feels sloppy:

maskIndx = np.c_[row_max_index-2, row_max_index-1, row_max_index, row_max_index+1, row_max_index+2,]%12
maskIndx #show windowed index
array([[ 7,  8,  9, 10, 11],
       [ 4,  5,  6,  7,  8],
       [ 8,  9, 10, 11,  0],
       [ 0,  1,  2,  3,  4],
       [10, 11,  0,  1,  2]])

results[np.meshgrid(np.arange(5), np.arange(5))[1], maskIndx] = 0 #uses array indexing 
results #show results
array([[ 33, 108, 208,  96,  96, 112,  18,   0,   0,   0,   0,   0],
       [ 51,   6,  39, 112,   0,   0,   0,   0,   0, 136,  42,  54],
       [  0, 176, 266, 162,  45,  70, 156, 198,   0,   0,   0,   0],
       [  0,   0,   0,   0,   0, 238, 156, 126, 165,  91, 144, 130],
       [  0,   0,   0,  51,  10, 128, 156, 272,  32,  98,   0,   0]])
next_max_index = results.argmax(1)
array([2, 9, 2, 5, 7])

Any ideas on doing this faster through indexing/windowing?

Answer

You can create a mask around the indices you compute for the max by taking an array of indices and subtracting the relevant max_indices, and then use the masked api to recompute the argmax:

import numpy as np

result = ... # Result here
row_max_index = results.argmax(axis=1, keepdims=True)

indices = np.arange(results.shape[1])
mask = np.abs(indices - row_max_index) <= 2

out = np.ma.array(results, mask=mask).argmax(axis=1)