Etant donné une matrice de numpy 2-D, comment conserver les N plus petits éléments de chaque rangée et le changement reste d'entre eux à 0(zéro).
Par exemple: N=3tableau d' entrée:
1 2 3 4 5
4 3 6 1 0
6 5 3 1 2
Production attendue:
1 2 3 0 0
0 3 0 1 0
0 0 3 1 2
Voici le code que je l'ai essayé et il fonctionne:
# 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
Toutefois, cette opération implique itérer sur chaque élément. Y at - il un moyen plus rapide de résoudre ce en utilisant np.argsort()ou toute autre fonction?
Approche # 1
Voici un avec l' np.argpartitionefficacité de la performance -
N = 3
newval = 0
np.put_along_axis(a,np.argpartition(a,N,axis=1)[:,N:],newval,axis=1)
Explication: Nous partitionner le tableau d'entrée pour obtenir des indices qui sont cloisonnés-partout pour l' kthargument np.argpartition. Donc, considérez cela comme essentiellement deux partitions, avec un premier pour les éléments de N plus petites le long de cet axe et l'autre pour le reste. Nous avons besoin de réinitialiser la deuxième partition, que nous sélectionnons avec [:,N:]et nous utilisons np.put_along_axispour faire la remise à zéro.
run échantillon -
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]])
Approche # 2
Voici une autre fois avec np.argpartition, mais couper juste le plus petit élément Nième par ligne, puis remettre à zéro tous supérieurs à lui. Ainsi, s'il y a des doublons pour le Nième élément le plus petit, nous gardons tous ceux avec cette méthode. Voici la mise en œuvre -
a[a>=a[np.arange(len(a)), np.argpartition(a,3,axis=1)[:,3],None]] = 0
Timings sur une version plus grande échelle -
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)
Cet article est collecté sur Internet, veuillez indiquer la source lors de la réimpression.
En cas d'infraction, veuillez [email protected] Supprimer.
laisse moi dire quelques mots