np.where() for conditional selection in NumPy

Introduction

We use np.where() to pick values from arrays based on a condition. It helps us choose data quickly without loops.

You want to replace values in a list based on a rule, like changing all negative numbers to zero.

You need to select elements from two arrays depending on a condition, like choosing prices based on a sale flag.

You want to find the positions of elements that meet a condition, like all scores above 90.

You want to create a new array that marks data points as 'pass' or 'fail' based on a threshold.

Syntax

NumPy

np.where(condition, x, y)

condition is a test that returns True or False for each element.

If condition is True, np.where() picks the value from x, else from y.

Examples

Replace values less or equal to 3 with 0, keep others.

NumPy

import numpy as np
arr = np.array([1, 2, 3, 4, 5])
result = np.where(arr > 3, arr, 0)
print(result)

Choose elements from a if greater than 15, else from b.

NumPy

a = np.array([10, 20, 30])
b = np.array([1, 2, 3])
cond = a > 15
result = np.where(cond, a, b)
print(result)

Find indices where elements are multiples of 10.

NumPy

arr = np.array([5, 10, 15, 20])
indices = np.where(arr % 10 == 0)
print(indices)

Sample Program

This program labels each temperature as 'Hot' if above 25 degrees, otherwise 'Cold'. It shows how np.where() helps assign categories based on a condition.

NumPy

import numpy as np

# Create an array of temperatures in Celsius
temps = np.array([22, 35, 18, 27, 30, 15, 40])

# We want to label temperatures above 25 as 'Hot' and others as 'Cold'
labels = np.where(temps > 25, 'Hot', 'Cold')

print('Temperatures:', temps)
print('Labels:', labels)

Important Notes

np.where() works element-wise on arrays, so it is very fast and efficient.

If you use np.where(condition) with only one argument, it returns the indices where the condition is True.

Summary

np.where() helps select values from two arrays based on a condition.

It can also find positions of elements that meet a condition.

This function is useful for quick, readable conditional data selection without loops.