Top K Documents Under a Latency Budget

Updated · By

Keep a min-heap of at most k scores while scanning once. Push until it holds k; after that, a score enters only if it beats the heap's smallest, which it replaces. Sort the k survivors at the end: O(n log k) time and O(k) space, against O(n log n) for a full sort.

hardPattern: HeapsFrom Devana's Google question set

Problem

A search query matches a large set of documents, each with a precomputed relevance score between 0 and 1. The results page needs only the best 20, and the latency budget does not allow sorting every match.

Write top_k_documents(scores, k), which returns the k highest scores in descending order without fully sorting the input, and state the complexity.

def top_k_documents(scores: list[float], k: int) -> list[float]:

This is one of Devana's own questions, written in the style of Google's interviews. It is not a question leaked from Google, and Devana is not affiliated with Google.

Constraints

  • Up to 10^5 matches per query; k is between 0 and 100 and usually much smaller than n.
  • Scores are between 0.0 and 1.0, and equal scores are interchangeable.
  • Return the scores in descending order.
  • If k exceeds n, return all of them; if k is 0, return an empty list.

Examples

Example 1

Input:  scores = [0.9, 0.1, 0.7, 0.5, 0.8], k = 3
Output: [0.9, 0.8, 0.7]

The three highest, best first.

Example 2

Input:  scores = [0.3], k = 5
Output: [0.3]

Asking for more than exist returns everything, rather than padding or failing.

Hints

Open one at a time, and try the problem again after each.

Hint 1

Sorting everything costs O(n log n). You only care about the order of k items.

Hint 2

Suppose you already hold the best k seen so far. What single number decides whether a new score gets in?

Hint 3

That number is the smallest of the k you hold. Which structure gives you the smallest item in O(1) and replaces it in O(log k)?

Approach: from brute force to optimal

Brute force: sort

Sort in descending order and take the first k. One line, often fast enough, but it orders all n items to use k of them.

Complexity: O(n log n) time, O(n) space for the sorted copy

A min-heap of size k

Scan once. While the heap holds fewer than k items, push. After that, compare each score with the heap's minimum and replace it only when the new score is larger. Most scores fail that comparison at O(1) cost. At the end the heap holds the top k; sort just those.

Complexity: O(n log k) time, O(k) space

Quickselect

Partition around a pivot until the k largest are on one side, then sort those. Linear on average, quadratic in the worst case without a careful pivot, and it needs the whole array in memory. The heap works on a stream, which matters for the follow-ups.

Complexity: O(n + k log k) on average

How to say it in the interview

State the trade explicitly: "sorting all n is O(n log n); I only need k, so a min-heap of size k gives O(n log k), and k is at most 100." Then say that the comparison with the heap's top is why most items cost almost nothing, and mention quickselect as the alternative you considered.

Solution

Python

import heapq

def top_k_documents(scores, k):
    if k <= 0:
        return []
    best = []                               # min-heap of the best k so far
    for score in scores:
        if len(best) < k:
            heapq.heappush(best, score)
        elif score > best[0]:               # beats the k-th best: swap it in
            heapq.heapreplace(best, score)
    return sorted(best, reverse=True)

JavaScript

class MinHeap {
  constructor() {
    this.items = [];
  }

  get size() {
    return this.items.length;
  }

  peek() {
    return this.items[0];
  }

  push(value) {
    const a = this.items;
    a.push(value);
    let i = a.length - 1;
    while (i > 0) {
      const parent = (i - 1) >> 1;
      if (a[parent] <= a[i]) break;
      [a[parent], a[i]] = [a[i], a[parent]];
      i = parent;
    }
  }

  replaceTop(value) {
    // Overwrite the smallest item and sift it down.
    const a = this.items;
    a[0] = value;
    let i = 0;
    for (;;) {
      const left = 2 * i + 1;
      const right = left + 1;
      let smallest = i;
      if (left < a.length && a[left] < a[smallest]) smallest = left;
      if (right < a.length && a[right] < a[smallest]) smallest = right;
      if (smallest === i) return;
      [a[smallest], a[i]] = [a[i], a[smallest]];
      i = smallest;
    }
  }
}

function topKDocuments(scores, k) {
  if (k <= 0) return [];
  const best = new MinHeap(); // the best k so far, smallest on top
  for (const score of scores) {
    if (best.size < k) best.push(score);
    else if (score > best.peek()) best.replaceTop(score); // beats the k-th best
  }
  return best.items.sort((a, b) => b - a);
}

Both versions are tested against the same cases, including every example on this page.

  • heapreplace pops and pushes in a single sift, which is cheaper than a separate pop and push.
  • JavaScript has no heap in its standard library, so the class above is the usual 40 lines. In an interview, say you would write it and ask whether you may assume one exists.
  • A score equal to the heap's top does not enter, which is fine because equal scores are interchangeable.

Follow-up questions

Scores are updated continuously by a ranking model. Does your structure still work?

Per query, yes: the heap is built fresh from a snapshot of the scores. A long-lived top k over changing scores needs a structure that can update an item's key, such as an indexed heap or a balanced tree.

The documents live on 50 shards. How do you get a correct global top 20?

Ask each shard for its own top 20 and merge. The global top 20 has to be inside that union, because a document outside its shard's top 20 already has 20 better documents on its own shard. That is 1,000 candidates, and one more heap gives the answer.

Can you stop early if you know the maximum possible remaining score?

Yes, if documents arrive ordered by an upper bound on their score. Once the next bound is below the heap's top, nothing left can enter, so stop.

What about k close to n?

Then the log k factor buys nothing. Keep a heap of the n minus k items to discard instead, or simply sort.

Common mistakes

  • Building a max-heap of all n items, which works but holds everything in memory and cannot run on a stream.
  • Keeping a max-heap of size k, whose top is the best item rather than the one to evict.
  • Forgetting k = 0 and k > n.
  • Sorting with JavaScript's default comparator, which compares numbers as strings.

Practice it in a mock interview

Reading a solution and explaining one under questioning are different skills. The button below starts this problem as a voice mock interview with Josh, Devana's AI interviewer, with an editor that runs your code.

Practice this problem in a mock interview

Practice this problem out loud

Josh, Devana's AI interviewer, runs this exact problem as a voice mock interview, asks follow-ups like the ones above, and scores you out of 100. The Free plan has 3 interviews a month, no card needed.

Practice this problem in a mock interview