Conditional Logic, Sorting & Set Logic

What you’ll learn

• Replace if/else loops with the vectorized np.where
• Sort arrays with np.sort (copy) vs arr.sort() (in place), and along an axis
• Recover a sort order with argsort instead of the sorted values
• Use set-style operations: np.unique, np.isin, intersect1d, union1d, setdiff1d
• Count and test conditions on whole arrays with any, all, and boolean sum

Conditional logic without a loop

Suppose you had two arrays and a boolean array telling you, position by position, which one to pick from. In plain Python you’d reach for a list comprehension:

python-ternary

xarr = [1.1, 1.2, 1.3, 1.4, 1.5]
yarr = [2.1, 2.2, 2.3, 2.4, 2.5]
cond = [True, False, True, True, False]
 
result = [(x if c else y) for x, y, c in zip(xarr, yarr, cond)]
print(result)

This works, but it’s slow on large arrays (every step runs in interpreted Python) and it doesn’t generalize to multidimensional data. np.where is the vectorized version of the ternary expression x if condition else y:

np-where

import numpy as np
 
xarr = np.array([1.1, 1.2, 1.3, 1.4, 1.5])
yarr = np.array([2.1, 2.2, 2.3, 2.4, 2.5])
cond = np.array([True, False, True, True, False])
 
result = np.where(cond, xarr, yarr)
print(result)

diagram How np.where chooses mermaid

  flowchart LR
  A["cond[i]"] --> D{"cond[i] is True?"}
  B["xarr[i]"] --> D
  C["yarr[i]"] --> D
  D -- "Yes" --> E["take xarr[i]"]
  D -- "No" --> F["take yarr[i]"]
  E --> G["result[i]"]
  F --> G

np.where walks cond, xarr, and yarr together and picks the matching value at each position.

The second and third arguments don’t have to be arrays

A common analytics pattern: replace every value that matches a condition with a constant, and leave the rest untouched. Either side of np.where can be a scalar, an array, or a mix:

where-scalar

import numpy as np
 
rng = np.random.default_rng(seed=1)
arr = rng.standard_normal((4, 4))
 
# Replace every positive value with 2, keep negatives as-is
capped = np.where(arr > 0, 2, arr)
 
# Replace positives with 2 and negatives with -2
signed = np.where(arr > 0, 2, -2)

where-clip-negatives

import numpy as np
 
arr = np.array([-3, 5, -1, 8, 0, -7])
cleaned = np.where(arr > 0, arr, 0)
print(cleaned)

sketch np.where picks between two arrays p5.js Pause Restart

Each cell in the result takes its value from xarr (amber) when cond is True, or from yarr (blue) when cond is False.

Sorting arrays

NumPy arrays sort much like Python lists — but the choice between arr.sort() and np.sort(arr) matters:

• arr.sort() sorts the array in place (mutates it, returns None).
• np.sort(arr) returns a sorted copy, leaving the original untouched — like Python’s built-in sorted().

sort-inplace-vs-copy

import numpy as np
 
arr = np.array([5, -10, 7, 1, 0, -3])
 
sorted_copy = np.sort(arr)   # arr is unchanged
print(sorted_copy)
print(arr)
 
arr.sort()                   # now arr itself is sorted
print(arr)

Sorting along an axis

For 2D arrays, pass axis to sort every row or every column independently:

sort-axis

import numpy as np
 
mat = np.array([
    [9, 2, 7],
    [1, 8, 3],
])
 
mat_sorted_by_col = np.sort(mat, axis=0)  # sort within each column
mat_sorted_by_row = np.sort(mat, axis=1)  # sort within each row
 
print(mat_sorted_by_col)
print(mat_sorted_by_row)

Getting the sort order with argsort

Sometimes you don’t want the sorted values — you want the indices that would put the array in order, so you can reorder other arrays the same way:

argsort

import numpy as np
 
scores = np.array([50, 10, 40, 20])
order = np.argsort(scores)
print(order)               # indices that sort scores ascending
print(scores[order])       # apply that order to get sorted values

Set logic: unique, membership, and set operations

NumPy has a small toolkit of set-style operations for 1D arrays.

np.unique returns the sorted, unique values in an array — like sorted(set(x)), but faster and returning an ndarray:

unique

import numpy as np
 
names = np.array(["Bob", "Will", "Joe", "Bob", "Will", "Joe", "Joe"])
print(np.unique(names))

np.isin (the modern name for the older np.in1d) tests membership of one array’s values in another, returning a boolean array:

isin

import numpy as np
 
values = np.array([6, 0, 0, 3, 2, 5, 6])
print(np.isin(values, [2, 3, 6]))

intersect1d, union1d, setdiff1d work like Python set operations, but stay in NumPy and return sorted arrays:

set-ops

import numpy as np
 
a = np.array([1, 2, 3, 4, 5])
b = np.array([3, 4, 5, 6, 7])
 
print(np.intersect1d(a, b))  # common to both
print(np.union1d(a, b))      # everything, no duplicates
print(np.setdiff1d(a, b))    # in a but not in b

Methods for boolean arrays

Boolean values act as 1 (True) and 0 (False) inside NumPy’s aggregation methods. That makes .sum() a quick way to count how many elements satisfy a condition:

bool-sum

import numpy as np
 
arr = np.array([1, -2, 3, -4, 5, 6, -7])
print((arr > 0).sum())   # count of positive values

.any() and .all() test a whole array at once — handy for questions like “did any random walk cross this line?” or “are all values positive?”:

any-all

import numpy as np
 
bools = np.array([False, False, True, False])
print(bools.any())   # True — at least one True
print(bools.all())   # False — not every value is True

A quick word on running totals

Not every array method collapses data to a single summary. cumsum and cumprod keep every intermediate step, and both accept axis on multidimensional arrays:

cumsum

import numpy as np
 
steps = np.array([1, -1, 1, 1, -1, 1, 1, -1])
print(steps.cumsum())   # running total after each step

You’ll use exactly this pattern — np.where to turn coin flips into +1/-1 steps, then cumsum to turn steps into a position — to build a random walk in the next page.

From the book

McKinney calls np.where “a vectorized version of the ternary expression x if condition else y.” The key habit to build: whenever you catch yourself writing a Python for loop with an if/else inside it to build a new array, ask whether np.where can replace it in one line.

Next

Continue to: Structured Arrays & Random Walks to combine np.where, cumsum, and boolean reductions into a full simulation.

🧪 Try It Yourself

Exercise 1 – Clip Negatives with np.where

pch.exerciseTag pch.exerciseTitle

Exercise 2 – Sort a Copy, Not the Original

pch.exerciseTag pch.exerciseTitle

Exercise 3 – Unique Values with np.unique

pch.exerciseTag pch.exerciseTitle