Skip to content
Python Central Hub

Conditional Logic, Sorting & Set Logic

What you’ll learn

  • Replace if/elseif/else loops with the vectorized np.wherenp.where
  • Sort arrays with np.sortnp.sort (copy) vs arr.sort()arr.sort() (in place), and along an axis
  • Recover a sort order with argsortargsort instead of the sorted values
  • Use set-style operations: np.uniquenp.unique, np.isinnp.isin, intersect1dintersect1d, union1dunion1d, setdiff1dsetdiff1d
  • Count and test conditions on whole arrays with anyany, allall, and boolean sumsum

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)
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.wherenp.where is the vectorized version of the ternary expression x if condition else yx 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)
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
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.wherenp.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-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)
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
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()arr.sort() and np.sort(arr)np.sort(arr) matters:

  • arr.sort()arr.sort() sorts the array in place (mutates it, returns NoneNone).
  • np.sort(arr)np.sort(arr) returns a sorted copy, leaving the original untouched — like Python’s built-in sorted()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)
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 axisaxis 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)
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
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.uniquenp.unique returns the sorted, unique values in an array — like sorted(set(x))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))
unique
import numpy as np
 
names = np.array(["Bob", "Will", "Joe", "Bob", "Will", "Joe", "Joe"])
print(np.unique(names))

np.isinnp.isin (the modern name for the older np.in1dnp.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]))
isin
import numpy as np
 
values = np.array([6, 0, 0, 3, 2, 5, 6])
print(np.isin(values, [2, 3, 6]))

intersect1dintersect1d, union1dunion1d, setdiff1dsetdiff1d 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
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 11 (TrueTrue) and 00 (FalseFalse) inside NumPy’s aggregation methods. That makes .sum().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
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().any() and .all().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
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. cumsumcumsum and cumprodcumprod keep every intermediate step, and both accept axisaxis 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
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.wherenp.where to turn coin flips into +1+1/-1-1 steps, then cumsumcumsum to turn steps into a position — to build a random walk in the next page.

Next

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

🧪 Try It Yourself

Exercise 1 – Clip Negatives with np.where

Exercise 2 – Sort a Copy, Not the Original

Exercise 3 – Unique Values with np.unique

If this helped you, consider buying me a coffee ☕

Buy me a coffee

Was this page helpful?

Let us know how we did