Numpy получение индексов строк двух последних элементов каждого столбца в маске

Question

У меня есть логическая маска в форме (M, N). Каждый столбец в маске может иметь различное количество True элементов, но гарантированно будет иметь по крайней мере два. Я хочу найти индекс строки последних двух таких элементов настолько эффективно, насколько это возможно.

Если бы я хотел только один элемент, я мог бы сделать что-то вроде (M - 1) - np.argmax(mask[::-1, :], axis=0). Однако это не поможет мне получить второй по порядку индекс.

Я нашел итеративное решение с использованием np.where или np.nonzero:

M = 4
N = 3
mask = np.array([
    [False, True, True],
    [True, False, True],
    [True, False, True],
    [False, True, False]
])
result = np.zeros((2, N), dtype=np.intp)
for col in range(N):
    result[:, col] = np.flatnonzero(mask[:, col])[-2:]

Это создает ожидаемое result:

array([[1, 0, 1],
       [2, 3, 2]], dtype=int64)

Я бы хотел избежать финального l oop. Существует ли разумно векторизованная форма вышеперечисленного? Я ищу конкретно две строки, которые всегда гарантированно существуют. Общее решение для произвольного количества элементов не требуется.

Answer 1

Argsort делает это -

In [9]: np.argsort(mask,axis=0,kind='stable')[-2:]
Out[9]: 
array([[1, 0, 1],
       [2, 3, 2]])

Другой с cumsum -

c = mask.cumsum(0)
out = np.where((mask & (c>=c[-1]-1)).T)[1].reshape(-1,2).T

Специально для ровно двух рядов, в одну сторону с argmax -

c = mask.copy()
idx = len(c)-c[::-1].argmax(0)-1
c[idx,np.arange(len(idx))] = 0
idx2 = len(c)-c[::-1].argmax(0)-1
out = np.vstack((idx2,idx))