user2162806
user2162806

Reputation: 133

finding number of values returned by numpy.where

I would like to find the number of places where numpy.where has evaluated as true. The following solution works, but is pretty ugly.

b = np.where(a < 5)
num = (b[0]).shape[0]

I'm coming from a language where I need to check if num > 0 before proceeding to do something with the resulting array. Is there a more elegant way of getting num, or a more Pythonic solution than finding num?

(For those familiar with IDL, I'm trying to replicate its simple b = where(a lt 5, num).)

Upvotes: 5

Views: 2578

Answers (3)

root
root

Reputation: 80456

In [1]: import numpy as np

In [2]: arr = np.arange(10)

In [3]: np.count_nonzero(arr < 5)
Out[3]: 5 

or

In [4]: np.sum(arr < 5)
Out[4]: 5

If you have to define b = np.where(arr < 5)[0] anyway, use len(b) or b.size ( len() seems to be a tiny bit faster, but they are pretty much the same in terms of performance).

Upvotes: 7

reptilicus
reptilicus

Reputation: 10417

Another way

In [14]: a = np.arange(100)
In [15]: np.where(a > 1000)[0].size > 0
Out[15]: False
In [16]: np.where(a > 10)[0].size > 0
Out[16]: True

Upvotes: 0

Schuh
Schuh

Reputation: 1095

This only works if the condition of the where statement implies that the result is not evaluated to False. But then it's easy:

import numpy as np
a = np.random.randint(1,11,100)
a_5 = a[np.where(a>5)]
a_10 = a[np.where(a>10)]
a_5.any() #True
a_10.any() #False

So if you wish you can try this before assigning, of course:

a[np.where(a>5)].any() #True
a[np.where(a>10)].any() #False

Upvotes: 0

Related Questions