Geek Answers Handbook

How to get maximum N values ​​for each row in numpy ndarray?

We know how to do this when N = 1

import numpy as np

m = np.arange(15).reshape(3, 5)
m[xrange(len(m)), m.argmax(axis=1)]    # array([ 4,  9, 14])

What is the best way to get the vertex N when N> 1? (say 5)

+4
source share
3 answers

Partial sorting using np.partitioncan be a lot cheaper than a full sort:

gen = np.random.RandomState(0)
x = gen.permutation(100)

# full sort
print(np.sort(x)[-10:])
# [90 91 92 93 94 95 96 97 98 99]

# partial sort such that the largest 10 items are in the last 10 indices
print(np.partition(x, -10)[-10:])
# [90 91 93 92 94 96 98 95 97 99]

If you need the largest N elements to sort, you can call np.sortin the last N elements of your partially sorted array:

print(np.sort(np.partition(x, -10)[-10:]))
# [90 91 92 93 94 95 96 97 98 99]

This can be much faster than a full sort over the entire array if your array is large enough.

axis= np.partition / np.sort:

y = np.repeat(np.arange(100)[None, :], 5, 0)
gen.shuffle(y.T)

# partial sort, followed by a full sort of the last 10 elements in each row
print(np.sort(np.partition(y, -10, axis=1)[:, -10:], axis=1))
# [[90 91 92 93 94 95 96 97 98 99]
#  [90 91 92 93 94 95 96 97 98 99]
#  [90 91 92 93 94 95 96 97 98 99]
#  [90 91 92 93 94 95 96 97 98 99]
#  [90 91 92 93 94 95 96 97 98 99]]

:

In [1]: %%timeit x = np.random.permutation(10000000)
   ...: np.sort(x)[-10:]
   ...: 
1 loop, best of 3: 958 ms per loop

In [2]: %%timeit x = np.random.permutation(10000000)
np.partition(x, -10)[-10:]
   ....: 
10 loops, best of 3: 41.3 ms per loop

In [3]: %%timeit x = np.random.permutation(10000000)
np.sort(np.partition(x, -10)[-10:])
   ....: 
10 loops, best of 3: 78.8 ms per loop
+3

- :

np.sort(m)[:,-N:]
+2

partition, sort, argsortEtc. accept axis parameter

Let me shuffle some values

In [161]: A=np.arange(24)

In [162]: np.random.shuffle(A)

In [163]: A=A.reshape(4,6)

In [164]: A
Out[164]: 
array([[ 1,  2,  4, 19, 12, 11],
       [20,  5, 13, 21, 22,  3],
       [10,  6, 16, 18, 17,  8],
       [23,  9,  7,  0, 14, 15]])

Section:

In [165]: A.partition(4,axis=1)

In [166]: A
Out[166]: 
array([[ 2,  1,  4, 11, 12, 19],
       [ 5,  3, 13, 20, 21, 22],
       [ 6,  8, 10, 16, 17, 18],
       [14,  7,  9,  0, 15, 23]])

The 4 smallest values ​​of each row are the first, 2 the largest last; slice to get an array of the 2 largest:

In [167]: A[:,-2:]
Out[167]: 
array([[12, 19],
       [21, 22],
       [17, 18],
       [15, 23]])

Sorting is probably slower, but on such a small array, it probably doesn't matter much. Plus it allows you to choose any N.

In [169]: A.sort(axis=1)

In [170]: A
Out[170]: 
array([[ 1,  2,  4, 11, 12, 19],
       [ 3,  5, 13, 20, 21, 22],
       [ 6,  8, 10, 16, 17, 18],
       [ 0,  7,  9, 14, 15, 23]])
+2
source

More articles:


All Articles