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:

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:

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)

  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

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:

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)

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

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().

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:

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:

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:

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:

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:

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:

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?":

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:

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.

Check your understanding

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

Press Run to see what your code prints.

Exercise 2 – Sort a Copy, Not the Original

Press Run to see what your code prints.

Exercise 3 – Unique Values with np.unique

Press Run to see what your code prints.