The NumPy where() function is used to get the indices of an array that matches a certain condition.
For example, suppose you want to get the indices of all elements greater than 10. Here’s how you do it:
import numpy as np
arr = np.array([5, 10, 15, 20, 25, 30])
indices = np.where(arr >10)
print(indices)
Output:
(array([2, 3, 4, 5]),)
You can add the indices to the original array to create a new array that fulfills the condition:
import numpy as np
arr = np.array([5, 10, 15, 20, 25, 30])
indices = np.where(arr >10)
new_arr = arr[indices]
print(new_arr)
Output:
[15 20 25 30]
Now that you’ve learned how to create an array using the where() function, let’s see how you can filter an array with multiple conditions next.
Using where() with multiple conditions
The where() function can also be used to filter an array with multiple conditions.
To use where() to filter an array with multiple conditions, you can use the and (&) or or (|) operator. You need to put each condition separately inside parentheses.
For example, suppose you want to filter the array to get elements that are greater than 5 but less than 20:
import numpy as np
arr = np.array([5, 10, 15, 20, 25, 30])
indices = np.where((arr > 5) & (arr < 20))
new_arr = arr[indices]
print(new_arr)
Output:
[10 15]
As you can see, the where() function effectively filters the arr variable with two conditions.
Next, let’s test an or condition. Suppose you want to add all elements greater than 20 or 5:
import numpy as np
arr = np.array([5, 10, 15, 20, 25, 30])
indices = np.where((arr > 20) | (arr == 5))
new_arr = arr[indices]
print(new_arr)
Output:
[ 5 25 30]
Here, only elements greater than 20 or exactly 5 will be added to the new array.
Using np.logical_and() and np.logical_or()
As an alternative to the and/or operator, you can also use the np.logical_and() or np.logical_or() function.
You need to call the function inside the np.where() function. Here’s an example of calling the logical_and() function:
import numpy as np
arr = np.array([5, 10, 15, 20, 25, 30])
indices = np.where(np.logical_and(arr > 5, arr < 20))
new_arr = arr[indices]
print(new_arr) # [10 15]
When you have multiple AND conditions, you need to place the filter conditions as arguments for the np.logical_and() function.
Here’s an alternative when using the np.logical_or() function:
import numpy as np
arr = np.array([5, 10, 15, 20, 25, 30])
indices = np.where(np.logical_or(arr > 20, arr == 5))
new_arr = arr[indices]
print(new_arr) # [ 5 25 30]
Conclusion
Now you’ve learned how to use the np.where() function with multiple conditions. You can use the and/or operator, or you can use the np.logical_and() or np.logical_or() function.
I hope this tutorial is helpful. Happy coding! 👍