K-th greatest value in every growing prefix: Frontier Interview's solution

Question: K-th greatest value in every growing prefix

Companies: Google, Netflix

Difficulty: Easy

Problem statement

You are given arr, a permutation of the integers 1 to n (each value appears exactly once), and an integer k with 1 <= k <= n.

Read arr from left to right. After reading each prefix of length i, for every i from k to n, report the k-th greatest value in that prefix. Return the n - k + 1 answers in order.

def kth_greatest_in_prefixes(arr: list[int], k: int) -> list[int]: ...

Constraints

1 <= n <= 10^5
arr is a permutation of 1..n
1 <= k <= n

Example 1

arr = [3, 1, 4, 2, 5], k = 2
output: [1, 3, 3, 4]

Prefix [3, 1]: the 2nd greatest value is 1. Add 4 to get [3, 1, 4]: the answer is 3. Add 2: still 3. Add 5: the top two are 5 and 4, so the answer is 4.

Example 2

arr = [4, 3, 2, 1], k = 1
output: [4, 4, 4, 4]

With k = 1 the answer is the running maximum.

Aim for better than sorting every prefix. Then say whether the fact that arr is a permutation lets you do better still.

Reconstructed from the reported question; details may differ from what interviewers ask.

Approach

You read the permutation from left to right. After each element from position k onward, you report the k-th greatest value seen so far. This is "Kth largest element in a stream" with a fixed input and one extra fact: the values are exactly 1..n.

Brute force. For each prefix, sort it and read index i - k. That costs O(n² log n). It is fine as a correctness oracle, but say right away that it repeats work.

Min-heap of size k (the expected answer). Keep the k greatest values seen so far in a min-heap. Its top is the k-th greatest.

Invariant: the heap holds exactly the k largest values of the current prefix (once the prefix has length k).
On a new value x: if x > heap[0], call heapreplace. Otherwise x can never be among the top k, so drop it.
This costs O(n log k) time and O(k) space. It works for any distinct or non-distinct numbers and for a true stream.

Permutation trick, O(n). Because the values are distinct and lie in 1..n, the answer only moves up, and it moves along the values present in the prefix:

Keep present[v] and the current answer ans. Exactly k present values are >= ans.
A new x < ans changes nothing.
A new x > ans makes k + 1 values >= ans. The new k-th greatest is the smallest present value above ans. Mark x first, then step ans up to the next v with present[v].
ans never decreases and never passes n, so all the stepping adds up to O(n) over the whole run.

Lead with the heap, because it is general. Then offer the O(n) version as the payoff for noticing "permutation of 1..n". Interviewers who say "Easy" usually want the heap, plus a clear reason why it is correct.

Code

import heapq
from typing import List


def kth_greatest_heap(arr: List[int], k: int) -> List[int]:
    """Answer for every prefix length i in [k, n], via a size-k min-heap."""
    if not 1 <= k <= len(arr):
        raise ValueError("k must be in [1, len(arr)]")
    heap = arr[:k]
    heapq.heapify(heap)
    out = [heap[0]]
    for x in arr[k:]:
        if x > heap[0]:
            heapq.heapreplace(heap, x)
        out.append(heap[0])
    return out


def kth_greatest_permutation(arr: List[int], k: int) -> List[int]:
    """O(n) version that relies on arr being a permutation of 1..n."""
    n = len(arr)
    if not 1 <= k <= n:
        raise ValueError("k must be in [1, len(arr)]")
    present = [False] * (n + 2)
    for x in arr[:k]:
        present[x] = True
    ans = min(arr[:k])          # exactly k present values are >= ans
    out = [ans]
    for x in arr[k:]:
        present[x] = True
        if x > ans:
            ans += 1            # move to the next present value above ans
            while not present[ans]:
                ans += 1
        out.append(ans)
    return out


def kth_greatest_brute(arr: List[int], k: int) -> List[int]:
    return [sorted(arr[:i], reverse=True)[k - 1] for i in range(k, len(arr) + 1)]


if __name__ == "__main__":
    import random

    assert kth_greatest_heap([3, 1, 4, 2, 5], 2) == [1, 3, 3, 4]
    assert kth_greatest_permutation([3, 1, 4, 2, 5], 2) == [1, 3, 3, 4]
    assert kth_greatest_permutation([1], 1) == [1]
    assert kth_greatest_permutation([1, 2, 3, 4], 1) == [1, 2, 3, 4]
    assert kth_greatest_permutation([4, 3, 2, 1], 1) == [4, 4, 4, 4]
    assert kth_greatest_permutation([2, 4, 1, 3], 4) == [1]
    for _ in range(500):
        n = random.randint(1, 30)
        p = random.sample(range(1, n + 1), n)
        k = random.randint(1, n)
        want = kth_greatest_brute(p, k)
        assert kth_greatest_heap(p, k) == want
        assert kth_greatest_permutation(p, k) == want
    print("ok")

Complexity

Brute force: O(n² log n) time, O(n) extra space per prefix.
Heap: O(k) to heapify the first k, then O(log k) per later element, so O(n log k) time and O(k) space. Each query is O(1) (heap[0]).
Permutation trick: O(n) time in total. ans rises at most n times over the whole run, and each other step is O(1). Space is O(n) for present.
Output: n - k + 1 values in every version.

Test cases

arr = [3, 1, 4, 2, 5], k = 2 gives [1, 3, 3, 4]. Prefix [3, 1] gives 1, adding 4 gives 3, adding 2 gives 3, adding 5 gives 4.
arr = [1], k = 1 gives [1] (smallest input).
arr = [1, 2, 3, 4], k = 1 gives [1, 2, 3, 4] (k = 1 is the running maximum).
arr = [4, 3, 2, 1], k = 1 gives [4, 4, 4, 4] (decreasing input, the answer never moves).
arr = [2, 4, 1, 3], k = 4 gives [1] (k = n, one prefix, the answer is its minimum).
arr = [5, 4, 3, 2, 1], k = 3 gives [3, 3, 3] (all later values are below the answer).
arr = [1, 2, 3, 4, 5], k = 3 gives [1, 2, 3] (every new value pushes the answer up by one).
k = 0 or k > n raises ValueError.
Randomized check: 500 random permutations, both fast versions match the brute force.

Follow-ups

True stream, unknown n, values not a permutation. Use the heap. The O(n) trick needs bounded, distinct values. With duplicates the heap still works if "k-th greatest" counts duplicates; for "k-th distinct" keep a set and only push new values.
k changes per query (answer any k for the current prefix). A heap of fixed size no longer fits. Use a Fenwick tree over values 1..n: insert with add(x, 1), then find the k-th greatest by walking down the tree to the (i - k + 1)-th smallest. Each operation is O(log n). For unbounded values, compress coordinates offline, or use an order-statistic tree such as sortedcontainers.SortedList.
Deletions or a sliding window of size w. The heap cannot delete cheaply. Use a Fenwick tree over values (add -1 on removal), a SortedList, or two heaps with lazy deletion. Each step is O(log n).
Huge n, data does not fit in memory. The heap needs only O(k) memory, so it streams. If k is also huge, keep a value histogram (Fenwick or count array) for bounded values, or shard the data, keep each shard's top k, and merge them.
Return the k greatest values themselves, sorted, for each prefix. Keep them in a sorted structure of size k (SortedList or a balanced BST). The cost is O(log k) per update, plus O(k) to emit each list.

Interviewer expectations

Mid-level (L4)

Gives the brute force, then reaches the size-k min-heap without hints.
Explains why the heap top is the k-th greatest, and why a value at or below the top is dropped.
Writes working code with the right output length (n - k + 1) and gets the edge cases right: k = 1, k = n, and a single element.

Senior (L5)

States the heap invariant and the O(n log k) / O(k) bounds clearly.
Notices that the input is a permutation and finds the monotone-answer O(n) method, with the amortized argument.
Tests against a brute-force oracle and handles bad k.

Staff (L6+)

Explains which assumptions each method needs (distinct values, bounded range, append-only) and picks one for the real setting.
Moves easily to variable k, deletions, and sliding windows (Fenwick tree, order-statistic tree, lazy deletion).
Discusses memory and streaming limits for very large inputs, and how to shard and merge top-k results.