Table of Contents



Introduction

You can search an array for a certain value and return either the value or the indices corresponding to the value.

The numpy.extract Function

The numpy.extract function returns the elements of an array that satisfy some condition.

Syntax

The numpy.extract function.

1numpy.extract(condition, arr)
Parameter Required? Default Value Description
condition ✔️ Yes NA An array whose nonzero or True entries indicate the elements of arr to extract.
arr ✔️ Yes NA Input array of the same size as condition.

As usual, we begin this section by first importing the NumPy package and creating a few random arrays using the numpy.random.choice function and seeding the PRNG so that results are reproducible by the reader.

1import numpy as np
2R = np.random.RandomState(17) #set a random seed
3a1 = R.choice(10, size=15, replace=True) # 1-D random array
4print(a1)
5a2 = R.choice(10, size=(4,5), replace=True) # 2-D random array
6print(a2)
[1 6 6 9 0 6 4 7 4 7 1 1 9 8 2]

[[3 6 6 9 9]
 [1 5 1 0 5]
 [6 6 2 6 9]
 [8 3 2 1 9]]

Note the argument replace=True which means that repetitions of values are allowed.

Example

Extract elements of a 1-D array satisfying a condition.

1np.extract(a1 > 6, a1) # elements > 6
array([9, 7, 7, 9, 8])

Typically, we also want to know the total number of elements within the array that satisfy the given condition. This can be done easily using the len function.

Example

Total number of elements of a 1-D array satisfying a condition.

1len(np.extract(a1 > 6, a1))
5

Example

How many elements of a 1-D array are even?

1print(np.extract(a1%2==0, a1))
2print(len(np.extract(a1%2==0, a1)))
[6 6 0 6 4 4 8 2]

8

In the above example, we made use of the modulo symbol a1%2==0 as the condition to extract the even numbers.

ADVERTISEMENT



The numpy.where Function

The numpy.where function returns the index or indices of the value(s) satisfying a specified condition. Further, it can also return elements chosen from a given set depending on the specified condition.

Syntax

The numpy.where function.

1numpy.where(condition, x=None, y=None)
Parameter Required? Default Value Description
condition ✔️ Yes NA Where True, yield x, otherwise yield y.
x, y ❌ No None Values from which to choose. x, y and condition need to be broadcastable to some shape.
Note

If x, y are not specified, the function will return the indices of the values satisfying the condition.

Returning Indices

We will first focus on using the numpy.where function to return the indices of the items of an array satisfying a specified condition.

Example

Indices of elements of a 1-D array satisfying a condition.

In this example, only indices of elements with value > 6 are extracted.

1np.where(a1 > 6) # tuple containing indices of elements > 6
(array([ 3,  7,  9, 12, 13], dtype=int64),)

The output is a tuple where the first element is the array of indices. To extract the indices, we write:

1print(np.where(a1 > 6)[0]) # indices of elements > 6
[ 3  7  9 12 13]

We can also combine numpy.where with numpy.extract to return both indices and values.

Example

Return indices and values of a 1-D array satisfying a condition.

1for (index, value) in zip(np.where(a1 > 6)[0], np.extract(a1>6, a1)):
2    print(f'{index}: {value}')
3: 9
7: 7
9: 7
12: 9
13: 8

where we have invoked the zip function to create a list of tuples.

👀 Review

For a more structured presentation, we may employ the pandas package.

1data = {'Index': np.where(a1 > 6)[0], 'Value': np.extract(a1>6, a1)} # dictionary
2import pandas as pd
3pd.DataFrame(data) # pandas dataframe
   Index  Value
0      3      9
1      7      7
2      9      7
3     12      9
4     13      8

Example

Find the values (and their indices) of a 1-D array which are even.

1cond = a1%2==0 # condition (even integers)
2data = {'Index': np.where(cond)[0], 'Value': np.extract(cond, a1)}
3pd.DataFrame(data)
   Index  Value
0      1      6
1      2      6
2      4      0
3      5      6
4      6      4
5      8      4
6     13      8
7     14      2

ADVERTISEMENT



Returning Modified Array

We now discuss using the numpy.where function to return a modified version of the input array based on a specified condition. Recall that numpy.where(condition, x, y) will return x if condition is True and y if condition is False.

Example

Modify a 1-D array based on whether the values are even or not.

In this example, an element becomes 1 if it is even and 0 otherwise.

1np.where(a1%2==0, 1, 0)  # even => 1; otherwise return 0.
array([0, 1, 1, 0, 1, 1, 1, 0, 1, 0, 0, 0, 0, 1, 1])

We can make use of the above output to count the number of even elements using the sum function.

Example

Total number of even elements in an array.

1sum(np.where(a1%2==0, 1, 0))  # sum number of even elements
8

The above can also be done using the numpy.extract function illustrated earlier.

Example

Modify a 2-D array based on whether the value is equal to 1.

1np.where(a2==1, a2*100, a2)  # return 100 if value=1; otherwise do nothing.
array([[  3,   6,   6,   9,   9],
       [100,   5, 100,   0,   5],
       [  6,   6,   2,   6,   9],
       [  8,   3,   2, 100,   9]])

The numpy.insert Function

The numpy.insert function inserts values along the given axis before the given indices and outputs a copy of the input array with the values inserted.

Syntax

The numpy.insert function.

1numpy.insert(arr, obj, values, axis=None)
Parameter Required? Default Value Description
arr ✔️ Yes NA Input array.
obj ✔️ Yes NA Object that defines the index or indices before which values is inserted.
values ✔️ Yes NA Values to insert into arr. In addition, values should be shaped so that arr[…,obj,…] = values is legal.
axis ❌ No None Axis along which to insert values. If axis is None then arr is flattened first.
Caution

Array values should be shaped so that arr[...,obj,...] = values is legal.

Example

Insert values into a 1-D array using numpy.insert.

1ar = np.arange(10) # input ordered array
2print(ar)
3ins = [3, 5, 7] # insert indices
4val = [91, 92, 93] # values to be inserted
5print(np.insert(ar, ins, val)) # new array
[0 1 2 3 4 5 6 7 8 9]
[ 0  1  2 91  3  4 92  5  6 93  7  8  9]

The indices array specified is [3, 5, 7] correspoding to the values array [91, 92, 93]. The values are inserted before the specified indices of the input array.

Example

Insert values into a 2-D array using numpy.insert with axis=1.

1print(a2)
2ins = [1, 3] # insert indices
3val = [91, 92] # values to be inserted
4print(np.insert(a2, ins, val, axis=1)) # new array
[[3 6 6 9 9]
 [1 5 1 0 5]
 [6 6 2 6 9]
 [8 3 2 1 9]]

[[ 3 91  6  6 92  9  9]
 [ 1 91  5  1 92  0  5]
 [ 6 91  6  2 92  6  9]
 [ 8 91  3  2 92  1  9]]

Example

Insert values into a 2-D array using numpy.insert with axis=0.

1print(a2)
2ins = [1, 3] # insert indices
3val = [91, 92] # values to be inserted
4print(np.insert(a2, ins, val, axis=0)) # new array
ValueError: shape mismatch: value array of shape (2,) could not be broadcast to indexing result of shape (2,5)

The above returns an error because the shape of val violates the condition that arr[…,ins,…] = val must be legal.

The way to do it would be:

Example

Insert values into a 2-D array using numpy.insert with axis=0.

1ins = [1, 3] # insert indices
2v1 = np.repeat(91,5)
3v2 = np.repeat(92,5)
4print(np.insert(a2, ins, [v1, v2], axis=0)) # new array
[[ 3  6  6  9  9]
 [91 91 91 91 91]
 [ 1  5  1  0  5]
 [ 6  6  2  6  9]
 [92 92 92 92 92]
 [ 8  3  2  1  9]]

ADVERTISEMENT



The numpy.searchsorted Function

The numpy.searchsorted function returns indices of an array where elements should be inserted to maintain order.

Syntax

The numpy.searchsorted function.

1numpy.searchsorted(a, v, side='left', sorter=None)
Parameter Required? Default Value Description
a ✔️ Yes NA Input array. If sorter is None, then it must be sorted in ascending order, otherwise sorter must be an array of indices that sort it.
v ✔️ Yes NA Values to be inserted into a.
side ❌ No 'left' If 'left', the index of the first suitable location found is given. If 'right', return the last such index. If there is no suitable index, return either 0 or N (where N is the length of a).
sorter ❌ No None Optional array of integer indices that sort array a into ascending order. They are typically the result of argsort.
Caution

If sorter is not specified, then a must be sorted in ascending order, otherwise sorter must be an array of indices that sort it.

Example

Get indices of a 1-D array where val should be inserted to maintain order using numpy.searchsorted.

1ar = np.arange(10) # already sorted
2print(ar)
3val = [3, 5, 20]
4ins = np.searchsorted(ar , val)
5print(ins)
[0 1 2 3 4 5 6 7 8 9]

array([ 3,  5, 10], dtype=int64)

The above returns the insertion indices. The following outputs the resulting array.

Example

Get sorted 1-D array after insertion using numpy.insert.

1np.insert(ar, ins, val)
array([ 0,  1,  2,  3,  3,  4,  5,  5,  6,  7,  8,  9, 20])

What if the input array is unsorted?

Example

Using numpy.searchsorted when input array is unsorted.

Step 1: Sort the array using ndarray.argsort.

1print(f'Random array: {a1}') # random array
2ind = a1.argsort() # get indices for sorting
3print(f'Sorted array: {a1[ind]}') # sorted array
Random array: [1 6 6 9 0 6 4 7 4 7 1 1 9 8 2]
Sorted array: [0 1 1 1 2 4 4 6 6 6 7 7 8 9 9]

Step 2: Get insertion indices using numpy.searchsorted.

1val = [3, 5, 20]  # values to be inserted
2print(f'Values to be inserted: {val}') # sorted array
3ins = np.searchsorted(a1[ind] , val) # insertion indices
4print(f'Insertion indices: {ins}')
Values to be inserted: [3, 5, 20]
Insertion indices: [ 5  7 15]

Step 3: Get resulting array after insertion using numpy.insert.

1fa = np.insert(a1[ind], ins, val) # resulting array after insertion
2print(f'Array after insertion: {fa}')
Array after insertion: [ 0  1  1  1  2  3  4  4  5  6  6  6  7  7  8  9  9 20]

Alternatively, we may also compute the insertion indices using numpy.searchsorted with sorterdefined from ndarray.argsort.

Example

Computing insertion indices using numpy.searchsorted with sorter.

1inss = np.searchsorted(a1 , val, sorter=ind) # ind comes from ndarray.argsort.
2print(f'Insertion indices: {inss}')
Insertion indices: [ 5  7 15]

With the insertion indices, we can once again generate the final inserted array using numpy.insert just like before.