Top k items per row
Medium12 min10 hidden tests
A search service scores m candidate documents for each of n queries, and the results page shows the k best per query. Write top_k_per_row(scores, k) that returns, for every row, the column indices of the k highest scores, best first. When two scores tie, the lower column index comes first, so the page looks the same on every reload.
Drag the scores, set k
k best columns, ties to the lower indexTop 3: columns [1, 3, 2], best first. Columns 1 and 3 tie at 0.9, so the lower index, 1, comes first.
Top 3 columns: 1, 3, 2.
Examples
The two 0.9 scores tie, so column 1 comes before column 3
top_k_per_row([[0.2, 0.9, 0.5, 0.9]], 3)
returns [[1, 3, 2]]
With k = 1 each query keeps its single best column
top_k_per_row([[3, 1, 2], [-1, -5, -2]], 1)
returns [[0], [0]]
Requirements
scoresis a 2-D list or array of shape (n, m), andkis an integer from 1 to m.- Return an integer NumPy array of shape (n, k).
- Each row lists column indices from the highest score down.
- Equal scores put the lower column index first.
The idea behind it
Ranking is the last step of most retrieval and recommendation systems. A model scores many candidates per query or per user, and only the top k reach the screen. The same operation picks the top-k tokens when a language model samples text.
np.argsort returns indices in the order that sorts the values from low to high. To rank from high to low you can negate the scores. Sort stability matters for ties: a stable sort keeps equal values in their original order, so the lower index stays ahead. Reversing an ascending sort looks equivalent but flips the order of ties too.
Deterministic tie-breaking matters in production. Without it the same query can show a different page after a restart, which confuses users and breaks offline tests that compare rankings. For very large m, np.argpartition finds the top k in linear time without sorting everything, and you then sort only those k.