np.where()

Boolean indexing filters values but loses their positions and the array's shape. Often you need either "where (at which index) did this happen?" or "replace values with X if the condition holds, else Y" — for every element at once. np.where does both.

Definition

np.where is a NumPy function with two forms. np.where(condition) returns the indexes of the elements where the condition is True, as a tuple of arrays (one array per dimension). np.where(condition, x, y) returns a new array of the same shape as the condition, taking the element from x where the condition is True and from y where it is False — a vectorized if / else.

• Tuple of arrays — even for a 1D array you get (array([...]),), so you usually take [0]. For 2D you get (row_indexes, col_indexes).
• Same shape — the 3-argument form never drops elements, so the table keeps its shape (unlike arr[mask]).
• x and y — can be arrays of the same shape, or scalars (a single number is broadcast to every position).
• Vectorized if/else — no Python loop; NumPy checks every element in compiled code.

Syntax

np.where(condition)            # -> tuple of index arrays
np.where(condition, x, y)      # -> array: x where True, y where False

Example — 1-argument form: find the positions

import numpy as np

virat_runs = np.array([72, 45, 88, 31, 64, 0, 103, 52, 19, 77, 41, 93, 28, 66])

result = np.where(virat_runs >= 50)
print(result)
idx = result[0]
print("Indexes     :", idx)
print("Match nums  :", idx + 1)
print("Scores      :", virat_runs[idx])
print("First century at index:", np.where(virat_runs >= 100)[0][0])

(array([ 0,  2,  4,  6,  7,  9, 11, 13]),)
Indexes     : [ 0  2  4  6  7  9 11 13]
Match nums  : [ 1  3  5  7  8 10 12 14]
Scores      : [ 72  88  64 103  52  77  93  66]
First century at index: 6

How it works: np.where(cond) is the same as np.nonzero(cond). It returns a tuple because arrays can have many dimensions; for 1D the tuple has one array, so [0] gets it.

Example — 1-argument form on a 2D array

scores = np.array([
    [72, 45, 88, 31, 64,  0, 103, 52],   # row 0: Virat
    [56,  0, 74, 88, 32, 45,  61, 17],   # row 1: Rohit
    [45, 32,  0, 67, 18, 54,  78, 22],   # row 2: Hardik
])
players = ["Virat", "Rohit", "Hardik"]

rows, cols = np.where(scores == 0)
print("rows:", rows, " cols:", cols)
for r, c in zip(rows, cols):
    print(f"{players[r]} made a duck in match {c + 1}")

rows: [0 1 2]  cols: [5 1 2]
Virat made a duck in match 6
Rohit made a duck in match 2
Hardik made a duck in match 3

How it works: for 2D, you get two arrays: the row index and the column index of every True cell. Zipping them gives (row, col) pairs.

Example — 3-argument form: vectorized if / else

virat_runs = np.array([72, 45, 88, 31, 64, 0, 103, 52, 19, 77, 41, 93, 28, 66])

labels = np.where(virat_runs >= 50, "Fifty+", "-")
print(labels)

capped = np.where(virat_runs > 80, 80, virat_runs)     # cap at 80
print("Capped :", capped)

bonus = np.where(virat_runs >= 50, virat_runs + 5, virat_runs)  # +5 bonus for 50+
print("Bonus  :", bonus)

['Fifty+' '-' 'Fifty+' '-' 'Fifty+' '-' 'Fifty+' 'Fifty+' '-' 'Fifty+' '-'
 'Fifty+' '-' 'Fifty+']
Capped : [72 45 80 31 64  0 80 52 19 77 41 80 28 66]
Bonus  : [ 77  45  93  31  69   0 108  57  19  82  41  98  28  71]

How it works: for every position NumPy asks "is the condition True here?" and takes x or y from that same position. Scalars (80, "-") are used everywhere; arrays (virat_runs + 5) are read position by position. The result keeps all 14 elements.

Example — keeps the 2D shape (from the course notebook)

scores = np.array([
    [72, 45, 88, 31, 64,  0, 103, 52],   # row 0: Virat
    [56,  0, 74, 88, 32, 45,  61, 17],   # row 1: Rohit
    [45, 32,  0, 67, 18, 54,  78, 22],   # row 2: Hardik
])

bonus_scores = np.where(scores >= 50, scores + 5, scores)
print(bonus_scores)
print("Only 50+, others zero:")
print(np.where(scores >= 50, scores, 0))

[[ 77  45  93  31  69   0 108  57]
 [ 61   0  79  93  32  45  66  17]
 [ 45  32   0  72  18  59  83  22]]
Only 50+, others zero:
[[ 72   0  88   0  64   0 103  52]
 [ 56   0  74  88   0   0  61   0]
 [  0   0   0  67   0  54  78   0]]

How it works: compare with scores[scores >= 50], which would give a flat list. np.where(..., scores, 0) keeps the 3×8 table and just zeroes the other cells — perfect when you still need to know which player and which match.

Tip: np.where is also the standard way to avoid divide-by-zero results, e.g. np.where(balls > 0, runs / balls * 100, 0) for strike rate — see Errors & Special Values in the Array Operations chapter.

Comparison — boolean indexing vs np.where

Common mistakes

• ❌ Forgetting that the 1-arg form returns a tuple: len(np.where(a > 50)) is 1, not the count. ✅ Use np.where(a > 50)[0] or (a > 50).sum().
• ❌ Giving only x without y: np.where(a > 50, a) → ValueError: either both or neither of x and y should be given.
• ❌ Using Python's if on an array (if scores > 50:) → ambiguous truth value error. ✅ np.where(scores > 50, ...).
• ❌ Expecting np.where to skip computing x where the condition is False. ✅ Both x and y are fully computed first, so np.where(b > 0, a / b, 0) still triggers a divide-by-zero warning (the result is correct). Suppress it with np.errstate.

Interview tip

Q: Difference between np.where(cond) and np.where(cond, x, y)? A: With one argument it returns the indexes where cond is True (a tuple, same as np.nonzero). With three it returns a new array of cond's shape, choosing from x where True and y where False — a vectorized if/else.