Как получить индексы двух верхних значений каждой строки в массиве 2-D numpy, но с указанной областью c исключается?

Question

У меня есть двумерный массив, например:

p = np.array([[21,2,3,1,12,13],
             [4,5,6,14,15,16],
             [7,8,9,17,18,19]])
b = np.argpartition(p, np.argmin(p, axis=1))[:, -2:]  
com = np.ones([3,6],dtype=np.int)
com[np.arange(com.shape[0])[:,None],b] = 0
print(com)

b - это индексы двух верхних значений каждой строки в p:

b = [[0 5]
    [4 5]
    [4 5]]

com - это np.ones матрица, того же размера, что и p, элемент с индексом, равным b, изменится на 0.

Итак, результат:

com = [[0 1 1 1 1 0]
       [1 1 1 1 0 0]
       [1 1 1 1 0 0]]

Теперь у меня есть еще одно ограничение:

p[0:2,0:2] = [[21  2]
              [4  5]]

Числа в этих областях [0: 2,0: 2] не должны учитываться, поэтому результат должен быть:

b = [[4 5]
     [4 5]
     [4 5]]
com = [[1 1 1 1 0 0]
       [1 1 1 1 0 0]
       [1 1 1 1 0 0]]

Как я могу это сделать? Стоит ли использовать маску или что-то подобное?

Заранее спасибо!

Answer 1

Просто установите значения в этих срезах на низкое значение, гарантируя, что они не будут среди двух самых больших, затем используйте argpartition:

out = np.copy(p)
out[0:2,0:2] = -np.inf
np.argpartition(out, [-2,-1])[:, -2:] 

array([[4, 5],
       [4, 5],
       [4, 5]])