Skip to content

← Performance step 5 of 7

Medium Primitives

Sort once

Jump to the task

To find how many values are smaller than a query, you can compare the query with every value and count the matches. Sorting the values separately for every query adds work without helping that comparison.

When many queries share the same values, sort once and find each query's insertion position in the sorted sequence:

import torch

values = torch.tensor([3.0, 1.0, 4.0, 1.0, 5.0])
queries = torch.tensor([1.0, 2.0, 4.0, 6.0])
sorted_values, _ = values.sort()
result = torch.searchsorted(sorted_values, queries)
print(result.tolist())  # [0, 2, 3, 5]

Choose the insertion side

The default searchsorted behavior finds the left insertion position, before any equal values. That index counts values strictly less than the query. In the example, the query 1.0 goes before both existing ones, so its answer is zero. Using side="right" would instead count values less than or equal to the query, which is a different task.

The search requires a sorted sequence, unless you supply an appropriate sorter. This exercise sorts the values explicitly and leaves the queries in their original order.

Why reuse the sorted data

Under the usual comparison-sort cost model, sorting n values once and using binary searches for m queries takes O(n log n + m log n) work. Repeating a sort for every query adds an O(m n log n) term. Moving the sort alone is not the whole improvement: replacing each full comparison scan with a search also changes the amount of work per query.

For very few queries or small inputs, a direct comparison may be competitive. Tensor layout, device and implementation affect actual timings. The point is to recognize reusable work and choose an algorithm suited to the workload.

Your task

def ranks(values: torch.Tensor, queries: torch.Tensor) -> torch.Tensor

The inputs are one-dimensional float32 tensors of finite values on the same device. Return an integer tensor containing, for each query in order, the number of values strictly less than it. Sort at most once and stay within four dispatched operations. Duplicate values count separately.

An empty query tensor should produce an empty result. If values is empty, every query has rank zero. The example above returns [0, 2, 3, 5].

The harness converts the result to a list under ranks and reports whether the dispatch count meets the budget under cheap. The reference uses a sort and a searchsorted dispatch. Dispatches are not kernel counts or timing measurements, so the name cheap does not establish a measured speedup.