
NumPy where - How to get the subarray of a NumPy array satisfying a condition
To get the subarray of an array satisfying the condition, numpy.where
can be used like sql where.
import numpy
a = numpy.array([3, 5, 7, 9, 11])
b = numpy.where(a > 5)
print(b) # (array([2, 3, 4]),)
print(type(b)) # <class 'tuple'>
where
returns a tuple and the first element is a NumPy array. It is an array of the original indexes satisfying the condition a > 5
. So b
is, to be exact, not a subarray of a
.
import numpy
a = numpy.array([3, 5, 7, 9, 11])
b = numpy.where(a > 5)
c = a[b]
print(b) # (array([2, 3, 4]),)
print(c) # [ 7 9 11]
The above code shows how to get the subset satisfying a > 5
. b
is an array of indexes so a[b]
is the subarray of a
satisfying a > 5
.
Comments
Powered by Markdown