NumPy Array Search

Table of Contents

• Introduction
• The numpy.extract Function
• The numpy.where Function
  • Returning Indices
  • Returning Modified Array
• The numpy.insert Function
• The numpy.searchsorted Function

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

1numpy.extract(condition, arr)

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

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

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

5

Example

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

1numpy.where(condition, x=None, y=None)

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

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

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.

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

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

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

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

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

1numpy.insert(arr, obj, values, axis=None)

Example

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

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

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

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

1numpy.searchsorted(a, v, side='left', sorter=None)

Example

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

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

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

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.