Shortest subarray with at least k distinct values: Frontier Interview's solution

Question: Shortest subarray with at least k distinct values

Companies: Google, Uber

Difficulty: Medium

Problem statement

You are given an integer array arr and an integer k. Call a contiguous subarray good if it contains at least k different values.

Return the length of the shortest good subarray. If no subarray is good, return -1.

Implement:

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

Constraints

1 <= len(arr) <= 10^5
1 <= arr[i] <= 10^5
1 <= k <= len(arr)

Example 1

arr = [4, 4, 7, 4, 9, 7], k = 3
output: 3

The window [4, 9, 7] holds three different values. No window of length 2 can.

Example 2

arr = [5, 5, 5], k = 2
output: -1

The array has only one distinct value, so no subarray is good.

Aim for a solution faster than checking every subarray.

Approach

A subarray is good when it holds at least k distinct values. We want the
shortest good subarray.

Brute force. For every start i, extend j to the right with a set of
seen values and stop at the first j where the set reaches size k. Record
j - i + 1. That is O(n^2) time, too slow for n = 10^5, but it shows the key
fact: for a fixed start, the answer is the first end that reaches k
distinct values, and that end never moves left as the start moves right.

Optimal: sliding window with a frequency map. "At least k distinct" is
monotone: widening a window can only keep or raise its distinct count, and
narrowing can only keep or lower it. So two pointers work.

Move right across the array. Add arr[right] to a count map; when a
count goes from 0 to 1, distinct grows by one.
While distinct >= k, the window [left, right] is good. Record its
length, then drop arr[left] and move left forward. When a count goes
to 0, distinct falls by one.
Return the best length, or -1 if no window was ever good.

Why it is correct: for each right, the inner loop moves left to the
largest start that still gives a good window ending at right, so it records
the shortest good window ending at right. The shortest good subarray ends
somewhere, so we see it. Each index enters and leaves the window once.

Invariants:

count[v] is the number of times v appears in arr[left..right].
distinct equals the number of keys with a positive count.
After the inner loop, distinct < k, so the window is never good between
iterations.

Two quick exits help: if k is larger than the number of distinct values in
the whole array, the answer is -1; if k == 1, the answer is 1.

A useful fact for follow-ups: the shortest good window has exactly k
distinct values, because removing one element lowers the count by at most one.

Code

from collections import defaultdict
from typing import Dict, List, Optional, Tuple


def shortest_good_subarray(arr: List[int], k: int) -> int:
    """Length of the shortest subarray with at least k distinct values, or -1."""
    span = shortest_good_window(arr, k)
    return -1 if span is None else span[1] - span[0] + 1


def shortest_good_window(arr: List[int], k: int) -> Optional[Tuple[int, int]]:
    """Inclusive (start, end) of the leftmost shortest good subarray, or None."""
    if k <= 0:
        return (0, -1)  # the empty window already qualifies
    if k > len(set(arr)):
        return None

    count: Dict[int, int] = defaultdict(int)
    distinct = 0
    left = 0
    best: Optional[Tuple[int, int]] = None

    for right, value in enumerate(arr):
        count[value] += 1
        if count[value] == 1:
            distinct += 1

        # Shrink while the window is still good; each step is a candidate.
        while distinct >= k:
            if best is None or right - left < best[1] - best[0]:
                best = (left, right)
            gone = arr[left]
            count[gone] -= 1
            if count[gone] == 0:
                distinct -= 1
            left += 1

    return best


def brute_force(arr: List[int], k: int) -> int:
    best = -1
    for i in range(len(arr)):
        seen = set()
        for j in range(i, len(arr)):
            seen.add(arr[j])
            if len(seen) >= k:
                if best == -1 or j - i + 1 < best:
                    best = j - i + 1
                break
    return best


if __name__ == "__main__":
    import random

    assert shortest_good_subarray([4, 4, 7, 4, 9, 7], 3) == 3   # [4, 9, 7]
    assert shortest_good_subarray([5, 5, 5], 2) == -1
    assert shortest_good_subarray([8], 1) == 1
    assert shortest_good_subarray([1, 2, 1, 1, 1, 3], 3) == 5
    assert shortest_good_window([6, 6, 2, 9, 2], 3) == (1, 3)

    for _ in range(2000):
        n = random.randint(1, 12)
        a = [random.randint(1, 5) for _ in range(n)]
        k = random.randint(1, n)
        assert shortest_good_subarray(a, k) == brute_force(a, k), (a, k)
    print("ok")

Complexity

Time: O(n). right visits each index once and left advances at most n
times in total, with O(1) average work per hash-map update.
Space: O(d), where d is the number of distinct values (at most
min(n, 10^5)). Because values are bounded by 10^5, a fixed-size list of
counts can replace the dict for speed.
The len(set(arr)) precheck is O(n) time and O(d) space and is optional.

Test cases

arr = [4, 4, 7, 4, 9, 7], k = 3 → 3 (window [4, 9, 7]).
arr = [1, 2, 1, 1, 1, 3], k = 3 → 5 ([2, 1, 1, 1, 3]; the only
3 is at the end, the nearest 2 is far back).
arr = [5, 5, 5], k = 2 → -1 (only one distinct value).
arr = [8], k = 1 → 1 (single element, smallest input).
arr = [3, 1, 2], k = 3 → 3 (the whole array is the only good window).
arr = [1, 2, 3, 4, 5], k = 1 → 1 (any single element works).
arr = [2, 2, 2, 2, 7], k = 2 → 2 (good window only at the tail).
arr = [1, 2] * 50000, k = 2 → 2 (large input, checks linear time).
arr of 10^5 distinct values, k = 10^5 → 100000.

Follow-ups

Return the subarray, not only its length. Track (left, right) with
the best length, as shortest_good_window does. Ties: use a strict < to
keep the leftmost one.
Count good subarrays (at least k distinct). Same window: after the
inner loop for each right, every start in [0, left - 1] gives a good
subarray ending at right, so add left to the total. O(n). For
exactly k (LeetCode 992), use atLeast(k) - atLeast(k + 1) or
atMost(k) - atMost(k - 1).
Shortest subarray with exactly k distinct. Same answer as "at least
k", because the shortest good window always has exactly k distinct values.
Say why rather than writing new code.
Many queries for different k. Run the window once per query: O(n * q).
For all k at once, note the answer is non-decreasing in k; for each
right, the window lengths for each threshold come from the last
occurrence of each distinct value seen so far. Sorting the last-seen
positions gives every k for that right, O(n * d) total, which is fine
when d is small. Mention the trade-off and pick by constraints.
Streaming input or a sliding time window. Keep the count map and
distinct as running state; process each arrival as a right step and
report the current best. If the data has timestamps and old items expire,
evict them from the left with the same count bookkeeping.

Interviewer expectations

Mid-level (L4)

States the O(n^2) idea, then reaches the two-pointer window with a count
map without heavy hints.
Keeps distinct correct on both add and remove (0 → 1 and 1 → 0).
Handles -1 and k = 1; tests a few small cases by hand.

Senior (L5)

Explains why two pointers are valid (monotonic "at least k" condition) and
why each window step finds the shortest window ending at right.
Gives O(n) time with an amortized argument and O(d) space; notes the
array-of-counts option from the value bound.
Solves the counting follow-up and the exactly-k version through the
at-least / at-most subtraction.

Staff (L6+)

Points out that the shortest good window has exactly k distinct values and
uses it to merge variants.
Discusses multi-k queries and streaming with expiry, and picks an approach
from the constraints.
Writes small, tested code (for example a brute-force cross-check) and
talks about tie rules and the empty-window edge case without prompting.