NumPy Searching, Sorting and Filtering

Lesson 11 of 17

NumPy can find where values are, put them in order, and keep only the elements that match a condition — all without loops.

Search with where()

Example

Python
import numpy as np

arr = np.array([4, 7, 4, 9, 4])
print(np.where(arr == 4))

Output

Plain Text
(array([0, 2, 4]),)

where() returns a tuple with the matching indexes. With three arguments it becomes a vectorized if/else:

Example

Python
import numpy as np

marks = np.array([35, 72, 48, 90])
print(np.where(marks >= 40, "pass", "fail"))

Output

Plain Text
['fail' 'pass' 'pass' 'pass']

Search a sorted array with searchsorted()

searchsorted() works on a sorted array and returns the index where a value would be inserted to keep the order:

Example

Python
import numpy as np

arr = np.array([10, 20, 30, 40])
print(np.searchsorted(arr, 25))

Output

Plain Text
2

Sort with np.sort()

Example

Python
import numpy as np

print(np.sort(np.array([3, 1, 2])))
print(np.sort(np.array(["banana", "cherry", "apple"])))
print(np.sort(np.array([[3, 2, 4], [5, 0, 1]])))

Output

Plain Text
[1 2 3]
['apple' 'banana' 'cherry']
[[2 3 4]
 [0 1 5]]

np.sort() returns a sorted copy and leaves the original unchanged. On a 2-D array it sorts each row. To get the order instead of the values, use np.argsort():

Example

Python
import numpy as np

prices = np.array([300, 120, 250])
print(np.argsort(prices))

Output

Plain Text
[1 2 0]

Filter with a boolean mask

Comparing an array with a value gives an array of True/False. Use it as an index to keep only the True positions:

Example

Python
import numpy as np

arr = np.array([41, 42, 43, 44, 45])
mask = arr > 42

print(mask)
print(arr[mask])

Output

Plain Text
[False False  True  True  True]
[43 44 45]

Combine conditions with & (and) and | (or), and put each condition in parentheses:

Example

Python
import numpy as np

arr = np.arange(1, 11)
print(arr[(arr > 3) & (arr % 2 == 0)])

Output

Plain Text
[ 4  6  8 10]