Как сохранить N наименьших элементов в данной строке массива numpy?

Question

Учитывая матрицу 2-D numpy, как сохранить N наименьших элементов в каждой строке и изменить остальные из них на 0 (ноль).

Например: N=3 Массив ввода :

1   2   3   4   5
4   3   6   1   0
6   5   3   1   2

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

1   2   3   0   0
0   3   0   1   0
0   0   3   1   2

Ниже приведен код, который я пробовал, и он работает:

# distance_matrix is the given 2D array
N=3
for i in range(distance_matrix.shape[0]):
    n_th_largest = np.sort(distance_matrix[i])[N]
    for j in range(distance_matrix.shape[1]):
        distance_matrix[i][j] = np.where(distance_matrix[i][j]<n_th_largest,distance_matrix[i][j],0)

# return distance_matrix

Однако эта операция включает в себя повторение каждый элемент. Есть ли более быстрый способ решить эту проблему, используя np.argsort() или любую другую функцию?

Answer 1

Подход № 1

Вот один с np.argpartition для повышения производительности -

N = 3
newval = 0
np.put_along_axis(a,np.argpartition(a,N,axis=1)[:,N:],newval,axis=1)

Объяснение: Мы разбиваем входной массив, чтобы получить индексы, которые разделены по для аргумента kth в np.argpartition. Итак, в основном рассмотрим это как два раздела, первый для наименьших N элементов вдоль этой оси, а другой для остальных. Нам нужно сбросить второй раздел, который мы выбираем с помощью [:,N:], и мы используем np.put_along_axis для выполнения сброса.

Пример выполнения -

In [144]: a # input array
Out[144]: 
array([[1, 2, 3, 4, 5],
       [4, 3, 6, 1, 0],
       [6, 5, 3, 1, 2]])

In [145]: np.put_along_axis(a,np.argpartition(a,3,axis=1)[:,3:],0,axis=1)

In [146]: a
Out[146]: 
array([[1, 2, 3, 0, 0],
       [0, 3, 0, 1, 0],
       [0, 0, 3, 1, 2]])

Approach # 2

Вот еще один вариант с np.argpartition, но просто нарезая N-й наименьший элемент на строку и затем сбрасывая все больше, чем он. Таким образом, если есть дубликаты для N-го наименьшего элемента, мы сохраним все те с этим методом. Вот реализация -

a[a>=a[np.arange(len(a)), np.argpartition(a,3,axis=1)[:,3],None]] = 0

Синхронизация в увеличенной версии -

In [184]: a = np.array([[1,2,3,4,5],[4,3,6,1,0],[6,5,3,1,2]])

In [185]: a = np.repeat(a,10000,axis=0)

In [186]: %timeit np.put_along_axis(a,np.argpartition(a,3,axis=1)[:,3:],0,axis=1)
1.78 ms ± 5.89 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [187]: a = np.array([[1,2,3,4,5],[4,3,6,1,0],[6,5,3,1,2]])

In [188]: a = np.repeat(a,10000,axis=0)

In [189]: %timeit a[a>=a[np.arange(len(a)), np.argpartition(a,3,axis=1)[:,3],None]] = 0
1.54 ms ± 54.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Answer 2

Если мы можем использовать pandas и иметь ваши данные в кадре данных, это однострочная строка с .apply(..., axis=1) в каждой строке:

df.apply(lambda row: row.nsmallest(3), axis=1).fillna(0).astype(int)

   0  1  2  3  4
0  1  2  3  0  0
1  0  3  0  1  0
2  0  0  3  1  2

Примечания:

• Если могут быть связи, nsmallest() имеет аргумент keep, указывающий, как их обрабатывать
• В этом случае вам нужен .astype(int), так как введение NaN вызывает dtype для принудительного плавания. Этого можно было бы избежать, если бы вы написали пользовательскую функцию для замены нулями.

А вот пример, который сделает ваш пример воспроизводимым:

import pandas as pd
from io import StringIO

dat = """1   2   3   4   5
4   3   6   1   0
6   5   3   1   2"""

df = pd.read_csv(StringIO(dat), sep='\s+', header=None)