Utsav T
Utsav T

Reputation: 1545

Python 2D NumPy array comprehension

I am new to NumPy. I have a 2-D NumPy array containing floating point values. I wish to get the index of those elements which are greater than 70 % of a certain value, say t ,in the entire matrix.

output = [(1,2),(4,7),(7,1)] meaning arr[1][2], arr[4][7] and arr[7][1] have values greater than 70% of t

Using 2 loops to get job done is a fairly uncomplicated way. What is the most Pythonic way of getting it done (list comprehension etc.) ? Please point out any duplicates. Thanks !

Upvotes: 2

Views: 831

Answers (1)

hpaulj
hpaulj

Reputation: 231395

An example:

In [76]: arr=np.arange(20, dtype=float).reshape(4,5)

In [77]: arr
Out[77]: 
array([[  0.,   1.,   2.,   3.,   4.],
       [  5.,   6.,   7.,   8.,   9.],
       [ 10.,  11.,  12.,  13.,  14.],
       [ 15.,  16.,  17.,  18.,  19.]])

A boolean index that can select values from the array

In [79]: arr>15
Out[79]: 
array([[False, False, False, False, False],
       [False, False, False, False, False],
       [False, False, False, False, False],
       [False,  True,  True,  True,  True]], dtype=bool)

In [80]: arr[arr>15]
Out[80]: array([ 16.,  17.,  18.,  19.])

Indexes where the condition is true, which can also be used to select elements

In [81]: I=np.nonzero(arr>15)

In [82]: I
Out[82]: (array([3, 3, 3, 3], dtype=int32), array([1, 2, 3, 4], dtype=int32))

In [83]: arr[I]
Out[83]: array([ 16.,  17.,  18.,  19.])

Or turn the index tuple into a list of pairs

In [84]: list(zip(*I))
Out[84]: [(3, 1), (3, 2), (3, 3), (3, 4)]

In [87]: [arr[j] for j in zip(*I)]
Out[87]: [16.0, 17.0, 18.0, 19.0]

Upvotes: 3

Related Questions