NumPy · Module 8 of 13 · 8 concepts · about 11 min

Sorting, Filtering & Top-K

Keeping the rows you want with a boolean mask, and finding the best few without paying for a full sort.

At work: Picking the ones that matter: top recommendations, riskiest transactions, the nearest stores, the readings over a limit.

Boolean masks

What it is

A comparison over an array gives a boolean array of the same shape; indexing with it returns a 1-D array of the elements where it is True.

Why it matters

It is the vectorised if: filtering, counting and conditional assignment are all the same mask expression, and mask.sum() counts the survivors before you commit to anything.

In the real world: Selecting only the transactions over ₹1,00,000 for a manual fraud review.

Visual example

Boolean indexing

A boolean mask is one True/False per element; indexing with it keeps the True positions and flattens the rest away.

x, with the mask

	x	x > 0
0	3	True
1	-1	False
2	4	True
3	-5	False

x[x > 0]

	x
0	3
1	4

What the mask gives you

x[mask] → [3, 4]

mask.sum() → 2

np.nonzero(mask)[0] → [0, 2]

x[~mask] → [-1, -5]

Why this matters

Every filter, count and conditional write in NumPy is this: build a boolean array, then use it. It runs in compiled code over the whole array at once.

Common misconception

Wrong idea: the result keeps the shape, with the False positions blanked out.

Correct idea: the result is 1-D and shorter — only the True elements, positions discarded. Use np.where(mask, x, 0) to keep the shape.

1Comparex > 0 gives a boolean array of the same shape.
2Indexx[mask] gathers the True positions into a new 1-D array.
3Or assignx[mask] = 0 writes to exactly those positions, in place.

x = np.array([3, -1, 4, -5])
mask = x > 0
mask                # array([ True, False,  True, False])
x[mask]             # array([3, 4])
mask.sum()          # 2 — booleans sum as 0/1
x[x < 0] = 0        # conditional assignment in place

Using and, or or not between masks. They call __bool__ on a whole array and raise 'truth value is ambiguous'. The element-wise operators are &, |, ~, and each comparison needs its own parentheses because & binds tighter than <.

np.where

What it is

np.where(condition, a, b) picks element-wise from a where the condition holds and from b where it does not, broadcasting all three. Called with one argument it returns the indices of the True elements.

Why it matters

It is the shape-preserving conditional: unlike a mask selection, the result has the same shape as the input, so it can go straight back into a column of a feature matrix.

In the real world: Labelling each order “free shipping” or “paid shipping” depending on whether it's over the threshold, for every order at once.

x = np.array([3.0, -1.0, 4.0])
np.where(x > 0, x, 0.0)          # array([3., 0., 4.]) — same shape
np.where(x > 0, 1, -1)           # a sign feature
np.where(x > 0)                  # (array([0, 2]),) — indices, as a tuple
np.where(count > 0, total / np.maximum(count, 1), 0.0)

np.where evaluates both branches before choosing. np.where(d != 0, n / d, 0) still computes n / 0 and still emits a divide warning (and a nan that is then discarded). Guard the arithmetic itself with np.divide(..., where=...).

nonzero, argwhere and flatnonzero

What it is

np.nonzero(mask) returns a tuple of index arrays, one per axis; np.argwhere(mask) returns the same information as an (n, ndim) array of coordinates; np.flatnonzero gives flat positions.

Why it matters

Sometimes you need *where*, not *what* — to index a parallel array, to report row numbers, or to compute a gap length between known values.

In the real world: Finding the positions of every sensor reading over the safety limit, to know which machines to check.

m = np.array([[False, True], [True, False]])
np.nonzero(m)        # (array([0, 1]), array([1, 0])) — usable as an index
np.argwhere(m)       # [[0, 1], [1, 0]] — one row per hit
np.flatnonzero(m)    # array([1, 2])
rows = np.unique(np.nonzero(np.isnan(X))[0])   # which rows have a nan

Using np.argwhere's output as an index. It is a 2-D array of coordinates, so x[np.argwhere(m)] indexes axis 0 with coordinate pairs and gives a strange shape. np.nonzero's tuple is the one that indexes.

sort and argsort

What it is

np.sort(x) returns the sorted values; np.argsort(x) returns the indices that would sort it, so x[np.argsort(x)] is np.sort(x). Both take axis= and default to the last axis.

Why it matters

argsort is the useful one: it lets you order *other* arrays by this one — documents by score, rows of a matrix by a column — which sorting the values alone cannot do.

In the real world: Ordering products by predicted demand, and keeping track of which product each number belongs to.

Visual example

argsort

argsort returns positions, not values — the permutation that would put the array in order.

scores

	score
0	0.2
1	0.9
2	0.5

x[np.argsort(x)]

sorted, ascending

	score
0	0.2
1	0.5
2	0.9

The permutation

argsort → [0, 2, 1]

0.2: from position 0

0.5: from position 2

0.9: from position 1

Why this matters

Because it hands back indices, one argsort can reorder every parallel array at once — scores, ids, labels — keeping the rows aligned.

Common misconception

Wrong idea: argsort returns the sorted values.

Correct idea: it returns where each sorted value came from. Use it as an index — on the array itself, or on any array in the same order.

1RankFind the order the values would take: smallest first.
2Record positions[0, 2, 1] — the origin of each sorted slot.
3GatherIndex any aligned array with it to reorder it identically.

scores = np.array([0.2, 0.9, 0.5])
order = np.argsort(scores)          # array([0, 2, 1]) — ascending
scores[order]                       # array([0.2, 0.5, 0.9])
doc_ids[np.argsort(scores)[::-1]]   # ids, best first
X[np.argsort(X[:, 0])]              # sort rows of a matrix by column 0

There is no descending= argument. Reverse with [::-1], which also reverses the tie order, or sort -scores — which is wrong for unsigned integers, where negation wraps.

argpartition and top-k

What it is

np.argpartition(x, kth) rearranges indices so that everything before position kth is no larger than everything after it, in linear time. The top k are then the last k entries — unordered.

Why it matters

Selection is O(n) while sorting is O(n log n). On 80 million candidates, where the last stage of every retrieval pipeline lives, that is the difference between a usable system and a timeout.

In the real world: Picking the 10 best-scoring products out of 5 million for a user's homepage, without sorting all 5 million.

Visual example

Top-k

Partition, then sort only the survivors: the boundary is guaranteed, the order inside each side is not.

scores

	score
0	3
1	9
2	1
3	7
4	5

argpartition(x, -2)

after partition at -2

	score
0	3
1	1
2	5
3	7
4	9

Two problems, two costs

partition all n: O(n)

top 2 = last 2 slots

sort those 2: O(k log k)

full sort would be O(n log n)

Why this matters

Retrieval needs the best few out of millions. Finding them is cheap; ordering all of them is not, and only the few need ordering.

Common misconception

Wrong idea: argpartition gives you the top k, in order.

Correct idea: it gives you which k, in no order. The only promise is that nothing before the boundary exceeds anything after it.

1PartitionOne linear pass puts the k largest past position -k.
2Slice[-k:] is the candidate set — unordered.
3Sort the kargsort the k scores and index back to rank them.

k = 3
candidates = np.argpartition(scores, -k)[-k:]   # the k best, any order
order = np.argsort(scores[candidates])[::-1]    # sort only k elements
top = candidates[order]                          # the k best, best first

Assuming the partitioned block is sorted. It is not, in any predictable way — argpartition guarantees the boundary and nothing else. If you need them ranked, sort the k survivors, not the n inputs.

Practise it: Top-K Retrieval Without Sorting Everything

Per-row top-k

What it is

argpartition and argsort take axis=, so a (n, m) score matrix gives an (n, k) index array. Apply it with np.take_along_axis, not with plain indexing.

Why it matters

Recommending k items to each of a million users is one call, and the naive per-row Python loop over argsort is two orders of magnitude slower.

In the real world: Finding the top 5 recommendations for each of 100,000 users at once.

k = 2
idx = np.argpartition(scores, -k, axis=1)[:, -k:]
topk = np.take_along_axis(scores, idx, axis=1)
order = np.argsort(topk, axis=1)[:, ::-1]
ranked = np.take_along_axis(idx, order, axis=1)   # item ids, best first per row

Indexing with scores[idx] instead of take_along_axis. That treats the (n, k) array as row positions and returns (n, k, m) — a silent explosion in memory rather than an error.

searchsorted

What it is

np.searchsorted(sorted_array, values) returns, for each value, the insertion position that keeps the array sorted — a binary search per value, vectorised. side='right' breaks ties the other way.

Why it matters

It is the tool for bucketing into arbitrary bin edges, for aligning irregular timestamps to a grid, and for as-of joins where the correct answer is 'the last known value at or before this time'.

In the real world: Placing each customer's spend into a price band (bronze, silver, gold) using a sorted list of band limits.

edges = np.array([0.0, 1.0, 5.0, 10.0])
np.searchsorted(edges, [0.5, 5.0, 7.0])         # array([1, 2, 3])
np.digitize([0.5, 5.0, 7.0], edges)              # the same idea, named for bins
prev = np.searchsorted(event_ts, query_ts, side='right') - 1   # as-of lookup

Calling it on an array that is not sorted. There is no check: you get positions from a binary search over unordered data, which are wrong in a way that looks like plausible noise.

lexsort and stable ties

What it is

np.lexsort((secondary, primary)) sorts by several keys, last key first — the opposite of what the argument order suggests. np.argsort(..., kind='stable') keeps the original order among equal elements.

Why it matters

Deterministic tie-breaking is what makes a ranking reproducible, and reproducible rankings are what make an evaluation metric comparable between runs.

In the real world: Ranking sales reps by revenue, and by name when revenue ties, so the leaderboard never reshuffles between refreshes.

order = np.lexsort((doc_ids, -scores))   # by score descending, then by id
np.argsort(scores, kind='stable')        # equal scores keep input order

Reading lexsort's tuple left to right. The *last* array is the primary key; getting it backwards produces a ranking sorted by the tie-break.