2

I have an array:

arr = [23,34,2,55,5,13,44,3]

and I want to save the n largest numbers with their indices.
for example, for n = 2, I want

[(55,3), (44,6)]

I couldn't find an easy way of doing that.. I only found how to get the n larget items using nlargest or n largest indices using argpartition.

kmario23
  • 57,311
  • 13
  • 161
  • 150

3 Answers3

3

Disclaimer: not a numpy solution - you can combine enumerate , sorted, generator comprehension and list slicing:

n = 3
s = list(sorted( ((v,i) for i,v in enumerate([23,34,2,55,5,13,44,3]) ),reverse = True))[:n]

print(s)

Output:

 [(55, 3), (44, 6), (34, 1)]

Doku:

The generator is used to flip index and value produced by enumerate, the slicing reduces the amount of results.

Patrick Artner
  • 50,409
  • 9
  • 43
  • 69
2

Here is a NumPy-based solution:

In [207]: arr
Out[207]: array([23, 34,  2, 55,  5, 13, 44,  3])

# sort array `arr` in increasing order
In [208]: arr_sorted = arr[arr.argsort()]

# slice largest 3 elements
In [209]: largest_3 = arr_sorted[-3:][::-1]

In [210]: largest_3
Out[210]: array([55, 44, 34])

# get the indices that would sort the array `arr` (in increasing order)
In [211]: arr_sorted_idx = arr.argsort()

# grab the largest 3 indices in accordance with [209]/[210]
In [212]: largest_3_idx = arr_sorted_idx[-3:][::-1]

In [213]: largest_3_idx
Out[213]: array([3, 6, 1])

largest_3_idx are the indices of largest 3 elements in the original array arr.


If at all you need the result as a list of tuples, then use:

In [214]: list(zip(largest_3.tolist(), largest_3_idx.tolist()))
Out[214]: [(55, 3), (44, 6), (34, 1)]
kmario23
  • 57,311
  • 13
  • 161
  • 150
2

Here's a shorten Numpy version:

n=3 #number of largest elements to get
a = np.array([23, 34,  2, 55,  5, 13, 44,  3])
idx = a.argsort()[:-n-1:-1] #Use argsort get sorted index and slice backwards
list(zip(a[idx], idx)) #zip and list for tuples

Output:

[(55, 3), (44, 6), (34, 1)]

Let's plot some timings:

import numpy
import perfplot

def sb_numpy(a,n):
    idx = a.argsort()[:-n-1:-1] #Use argsort get sorted index and slice backwards
    return list(zip(a[idx], idx))

def pa_pyth(a,n):
    return list(sorted( ((v,i) for i,v in enumerate(a) ),reverse = True))[:n]


perfplot.show(
    setup=lambda n: numpy.random.randint(0,10e7, n),
    kernels=[
        lambda a: sb_numpy(a,3),
        lambda a: pa_pyth(a,3)
        ],
    labels=['sb_numpy', 'pa_pyth'],
    n_range=[2**k for k in range(15)],
    xlabel='N'
    )

Output:

enter image description here

Scott Boston
  • 147,308
  • 15
  • 139
  • 187