In this article we will discuss about numpy.where() and show it’s usage with examples.

Python’s numpy module provide a function to select elements based on condition i.e.

numpy.where(condition[, x, y])

Argument:
condition : A condional expression that returns a Numpy array of bool
x, y : Arrays (Optional i.e. either both are passed or not passed)

  • If all arguments –> condition , x & y are passed in numpy.where() then it will return elements selected from x & y depending on values in bool array yielded by condition. All 3 arrays must be of same size.
  • If x & y arguments are not passed and only condition argument is passed then it returns a tuple of arrays (one for each axis) containing  the indices of the elements that are True in bool numpy array returned by condition.

Let’s understand by examples,

Using numpy.where() with all arguments

Let’s pass a bool array and two other arrays to nump.where() i.e.

# call numpy.where() with 3 arguments
result = np.where([True, False, False], [1, 2, 4], [7, 8, 9])

Contents of the numpy array returned are,
[1 8 9]

numpy.where() iterates over the bool array and for every True it yields corresponding element array x and for every False it yields corresponding element from array y. So, basically it returns an array of elements from x where condition is True, and elements from y elsewhere.

Now instead of bool array as first argument we can pass a condition too that will return a bool numpy array i.e.

# Create a Numpy array from a list
arr = np.array([11, 12, 13])

# call numpy.where() with 3 arguments
result = np.where(arr < 12, [1, 2, 4], [7, 8, 9])

print('Contents of Numpy array returned: ', result, sep='\n')

output:
Contents of Numpy array returned: 
[1 8 9]

arr < 12 returned a bool array that was then used as first argument in numpy.where().

Using numpy.where() with condition argument only

Let’s create a new numpy array,

# Create a numpy array from list
arr = np.array([11, 12, 13, 14, 15, 16, 17, 15, 11, 12, 14, 15, 16, 17])

We can pass the condition argument only in numpy.where() i.e.
# pass condition expression only
result = np.where((arr > 12) & (arr < 16))

In this case condition expression returns a bool numpy array, which is eventually passed to numpy.where(). Which returns a tuple of arrays a new array containing the indices where bool array has True in it i.e.
(array([ 2,  3,  4,  7, 10, 11], dtype=int32),)

As our numpy array had only one axis, so returned tuple will contain only one array of indices of elements whose value is between 12 & 16.

Some Important points:

  • We can either pass all the 3 arguments or pass one condition argument only. there can not be two arguments in case of numpy.where().
  • If we are passing all 3 arguments to numpy.where(). Then all the 3 numpy arrays must be of same length otherwise it will raise following error,
    • ValueError: operands could not be broadcast together with shapes

Complete example is as follows:

import numpy as np

def main():
    print('*** Use numpy.where() with All arguments ***')

    # call numpy.where() with 3 arguments
    result = np.where([True, False, False], [1, 2, 4], [7, 8, 9])

    print('Contents of Numpy array returned: ', result, sep='\n')

    # Create a Numpy array from a list
    arr = np.array([11, 12, 13])

    # call numpy.where() with 3 arguments
    result = np.where(arr < 12, [1, 2, 4], [7, 8, 9])

    print('Contents of Numpy array returned: ', result, sep='\n')

    print('*** Use numpy.where() condition argument only ***')

    # Create a numpy array from list
    arr = np.array([11, 12, 13, 14, 15, 16, 17, 15, 11, 12, 14, 15, 16, 17])

    # pass condition expression only
    result = np.where((arr > 12) & (arr < 16))

    if len(result) > 0:
        print('Elements in given range exists at following Index :', result)
    else:
        print('Element Not Found')


if __name__ == '__main__':
    main()


Output:
*** Use numpy.where() with All arguments ***
Contents of Numpy array returned: 
[1 8 9]
Contents of Numpy array returned: 
[1 8 9]
*** Use numpy.where() condition argument only ***
Elements in given range exists at following Index : (array([ 2,  3,  4,  7, 10, 11], dtype=int32),)

 

Subscribe with us to join a list of 2000+ Programmers for weekly newsletter.