Exactly k stops

A campus shuttle logs the stop it serves every time it opens its doors. The planner wants to know how much of the day is spent circling a small set of stops.

The problem

The log is a list of stop ids in the order they were served, with repeats — the shuttle goes back and forth. A stretch is a contiguous run of log entries, identified by where it starts and ends, so two runs over the same stops at different times count separately.

Count the stretches that touch exactly k distinct stops.

Input. log — a list of integer stop ids, possibly empty. k — a positive integer.

Output. The number of contiguous stretches containing exactly k distinct stop ids.

Example.

log = [4, 1, 4, 2, 1], k = 2   ->  5

The five are entries 0-1, 0-2, 1-2, 2-3 and 3-4. Every longer stretch drags in a third stop, and every single entry touches only one.

A second example, where the shuttle never leaves one stop:

log = [3, 3, 3, 3], k = 1   ->  10
log = [3, 3, 3, 3], k = 2   ->  0

All ten stretches touch exactly one stop, so k = 1 collects every one of them and k = 2 collects none. A method that reports one stretch per ending position would say 4 here.

Constraints.

• 0 <= len(log) <= 10^5
• 1 <= stop id <= 10^5
• 1 <= k <= 10^5

Hints

Approach

Brute force

Take every start, extend to every end, keeping a set of the stops seen. That is n * (n + 1) / 2 stretches, each costing a set insert — about 5 * 10^9 operations.

The insight

Windows with "at most k distinct" can be counted in one pass, and exactly k is the difference of two such counts: atMost(k) - atMost(k - 1).

The subtraction works because the two sets are nested: every stretch with at most k - 1 distinct stops also has at most k, so it is removed exactly once and what remains has a distinct count of exactly k. Counting "at most k" is cheap because the property is monotone — shrinking a window can never raise its distinct count — so one left pointer moving forward is enough.

Algorithm

1. Write a helper at_most(limit) that returns 0 when limit is 0.
2. Walk right across the log, incrementing a count map for the entry.
3. While the map holds more than limit keys, drop the entry at left, deleting the key when its count reaches 0, and advance left.
4. Add right - left + 1 to the total: that is the number of stretches ending at right that are legal.
5. Return at_most(k) - at_most(k - 1).

Complexity

Time O(n) — each helper moves both pointers forward at most n times, and the helper runs twice. Space O(n) for the count map, which holds at most one key per distinct stop.

Solution

"""Exactly k stops — count windows with at most k distinct, minus at most k-1."""

from collections import defaultdict


def at_most(log, limit):
    """Number of contiguous stretches holding at most `limit` distinct codes."""
    if limit <= 0:
        return 0
    counts = defaultdict(int)
    total = 0
    left = 0
    for right, code in enumerate(log):
        counts[code] += 1
        while len(counts) > limit:          # invariant: [left, right] is legal after this loop
            counts[log[left]] -= 1
            if counts[log[left]] == 0:
                del counts[log[left]]       # a zero count is not a distinct stop
            left += 1
        total += right - left + 1           # every stretch ending at right, starting at or after left
    return total


def solve(log, k):
    return at_most(log, k) - at_most(log, k - 1)

TESTS = [
    (([4, 1, 4, 2, 1], 2), 5),
    (([3, 3, 3, 3], 1), 10),
    (([3, 3, 3, 3], 2), 0),
    (([7, 5, 7, 5, 7], 2), 10),
    (([1, 2], 3), 0),
    (([], 1), 0),
    (([9], 1), 1),
]

Pitfalls

• Leaving zero counts in the map. len(counts) then never falls, the window keeps shrinking, and [4, 1, 4, 2, 1] with k = 2 reports 2 instead of 5. The key has to be deleted when its count hits 0.
• Calling at_most(0) without a guard. With k = 1 the second call gets limit 0, and the shrink loop walks left past right and off the log. Return 0 immediately instead — no non-empty stretch touches zero stops.
• Counting one stretch per ending position. The whole band of valid starts counts, which is why the total grows by right - left + 1 and not by 1. On [3, 3, 3, 3] with k = 1 that mistake gives 4 rather than 10.

Variants

• Two-ink run — the same at-most window with k fixed at 2, measuring the longest one instead of counting them.
• Sprint intervals — exactly k again, but over a yes/no property, where prefix counts beat the subtraction.