In Python, we can use the numpy.where() function to select elements from a numpy array, based on a condition.
Not only that, but we can perform some operations on those elements if the condition is satisfied.
Let’s look at how we can use this function, using some illustrative examples!
Table of Contents
Syntax of Python numpy.where()
This function accepts a numpy-like array (ex. a NumPy array of integers/booleans).
It returns a new numpy array, after filtering based on a condition, which is a numpy-like array of boolean values.
condition can take the value of
array([[True, True, True]]), which is a numpy-like boolean array. (By default, NumPy only supports numeric values, but we can cast them to
For example, if
array([[True, True, False]]), and our array is
a = ndarray([[1, 2, 3]]), on applying a condition to array (
a[:, condition]), we will get the array
import numpy as np a = np.arange(10) print(a[a <= 2]) # Will only capture elements <= 2 and ignore others
array([0 1 2])
NOTE: The same condition condition can also be represented as a <= 2. This is the recommended format for the condition array, as it is very tedious writing it as a boolean array
But what if we want to preserve the dimension of the result, and not lose out on elements from our original array? We can use numpy.where() for this.
numpy.where(condition [, x, y])
We have two more parameters
y. What are those?
Basically, what this says is that if
condition holds true for some element in our array, the new array will choose elements from
Otherwise, if it’s false, elements from
y will be taken.
With that, our final output array will be an array with elements from
condition = True, and elements from
condition = False.
Note that although
y are optional, if you specify
x, you MUST also specify
y. This is because, in this case, the output array shape must be the same as the input array.
NOTE: The same logic applies for both single and multi-dimensional arrays too. In both cases, we filter based on the condition. Also remember that the shapes of
condition are broadcasted together.
Now, let us look at some examples, to understand this function properly.
Using Python numpy.where()
Suppose we want to take only positive elements from a numpy array and set all negative elements to 0, let’s write the code using
1. Replace Elements with numpy.where()
We’ll use a 2 dimensional random array here, and only output the positive elements.
import numpy as np # Random initialization of a (2D array) a = np.random.randn(2, 3) print(a) # b will be all elements of a whenever the condition holds true (i.e only positive elements) # Otherwise, set it as 0 b = np.where(a > 0, a, 0) print(b)
[[-1.06455975 0.94589166 -1.94987123] [-1.72083344 -0.69813711 1.05448464]] [[0. 0.94589166 0. ] [0. 0. 1.05448464]]
As you can see, only the positive elements are now retained!
2. Using numpy.where() with only a condition
There may be some confusion regarding the above code, as some of you may think that the more intuitive way would be to simply write the condition like this:
import random import numpy as np a = np.random.randn(2, 3) b = np.where(a > 0) print(b)
If you now try running the above code, with this change, you’ll get an output like this:
(array([0, 1]), array([2, 1]))
If you observe closely,
b is now a tuple of numpy arrays. And each array is the location of a positive element. What does this mean?
Whenever we provide just a condition, this function is actually equivalent to
In our example,
np.asarray(a > 0) will return a boolean-like array after applying the condition, and
np.nonzero(arr_like) will return the indices of the non-zero elements of
arr_like. (Refer to this link)
So, we’ll now look at a simpler example, that shows us how flexible we can be with numpy!
import numpy as np a = np.arange(10) b = np.where(a < 5, a, a * 10) print(a) print(b)
[0 1 2 3 4 5 6 7 8 9] [ 0 1 2 3 4 50 60 70 80 90]
Here, the condition is
a < 5, which will be the numpy-like array
[True True True True True False False False False False],
x is the array a, and
y is the array a * 10. So, we choose from an only if a < 5, and from a * 10, if a > 5.
So, this transforms all elements >= 5, by multiplication with 10. This is what we get indeed!
Broadcasting with numpy.where()
If we provide all of
y arrays, numpy will broadcast them together.
import numpy as np a = np.arange(12).reshape(3, 4) b = np.arange(4).reshape(1, 4) print(a) print(b) # Broadcasts (a < 5, a, and b * 10) # of shape (3, 4), (3, 4) and (1, 4) c = np.where(a < 5, a, b * 10) print(c)
[[ 0 1 2 3] [ 4 5 6 7] [ 8 9 10 11]] [[0 1 2 3]] [[ 0 1 2 3] [ 4 10 20 30] [ 0 10 20 30]]
Again, here, the output is selected based on the condition, so all elements, but here,
b is broadcasted to the shape of
a. (One of its dimensions has only one element, so there will be no errors during broadcasting)
b will now become
[[0 1 2 3] [0 1 2 3] [0 1 2 3]], and now, we can select elements even from this broadcasted array.
So the shape of the output is the same as the shape of
In this article, we learned about how we can use the Python numpy.where() function to select arrays based on another condition array.
- SciPy Documentation on Python numpy.where() function