Лучшие N значений в 2d массиве с дубликатами для маскировки - PullRequest
Новая
       20

Лучшие N значений в 2d массиве с дубликатами для маскировки

0 голосов
/ 30 апреля 2020

У меня есть 2d numpy массив: 

arr = np.array([[0.1, 0.1, 0.3, 0.4, 0.5], 
                [0.06, 0.1, 0.1, 0.1, 0.01], 
                [0.24, 0.24, 0.24, 0.24, 0.24], 
                [0.2, 0.25, 0.3, 0.12, 0.02]])
print (arr)
[[0.1  0.1  0.3  0.4  0.5 ]
 [0.06 0.1  0.1  0.1  0.01]
 [0.24 0.24 0.24 0.24 0.24]
 [0.2  0.25 0.3  0.12 0.02]]

Я хочу фильтровать верхние значения N, поэтому я использую argsort: 

N = 2
arr1 = np.argsort(-arr, kind='mergesort') < N
print (arr1)
[[False False False  True  True]
 [ True False False  True False] <- first top 2 are duplicates
 [ True  True False False False]
 [False  True  True False False]]

Это работает хорошо, по крайней мере, нет верхние дубликаты, как для строки 2.

Ожидаемый результат: 

print (arr1)
[[False False False  True  True]
 [False  True  True False False]
 [ True  True False False False]
 [False  True  True False False]]

Возможно ли более быстрый способ справиться с этим?

1 Ответ

1 голос
/ 30 апреля 2020

Срез, чтобы получить эти верхние N индексы и использовать их для создания окончательной маски - 

idx = np.argsort(-arr, kind='mergesort')[:,:N]
mask = np.zeros(arr.shape, dtype=bool)
np.put_along_axis(mask, idx, True, axis=-1)
...