Conditional Selection with np.where() & np.select()
Conditional logic is the backbone of data engineering and scientific pipelines. Often, you need to assign values depending on whether conditions are met—similar to Excel's IF() function or SQL's CASE WHEN ... THEN ... ELSE statements.
NumPy provides two workhorses for conditional operations:
1
np.where(): Vectorized ternary operator if-else and index extraction.
2
np.select(): Scalable multi-condition decision tables for complex business logic.
1. Vectorized If-Else with np.where(condition, x, y)
The ternary syntax of np.where takes three arguments: np.where(condition, value_if_true, value_if_false)
Both value_if_true and value_if_false can be scalar constants or entire arrays of matching shape:
When called with only one argument (the condition), np.where() acts as an index finder, returning a tuple of coordinate arrays where the condition evaluates to True:
1voltages = np.array([1.2,3.8,5.5,0.4,4.9,5.1])
2
3# Locate indices where voltage exceeds 4.5V
4spike_indices = np.where(voltages >4.5)
5print("Indices of voltage spikes:", spike_indices)
6print("Actual values:", voltages[spike_indices])
Output:
1Indices of voltage spikes: (array([2, 4, 5], dtype=int64),)
2Actual values: [5.5 4.9 5.1]
In 2D Matrices:
In a 2D array, np.where() returns a tuple of (row_indices, col_indices):
1grid = np.array([
2[10,95,20],
3[85,30,90]
4])
5
6rows, cols = np.where(grid >80)
7print("Rows matching:", rows)
8print("Cols matching:", cols)
9
10for r, c inzip(rows, cols):
11print(f"Match found at ({r}, {c}) with value {grid[r, c]}")
3. Multi-Condition Logic with np.select()
When your logic involves more than two outcomes (e.g. grading scale A, B, C, D, F), nesting np.where() calls quickly turns into unreadable spaghetti code:
1# Bad practice: Nested np.where is unreadable
2# np.where(c1, a, np.where(c2, b, np.where(c3, c, d)))
Instead, use np.select(condlist, choicelist, default=default_value):