The Waiting Game

A medium Python interview practice problem on DataDriven. Write the Python and run it against test cases, with instant feedback.

Domain
Python
Difficulty
medium
Seniority
L4
Asked in
Citadel

Problem

A weather station logs one high-temperature reading per day, in recording order, and you get them as the list temps. For each day, return how many days you'd wait until a strictly warmer reading appears later in the log, using 0 when no later day is warmer.

Patience has a price - and a count.

Example 1

Input:

temps = [73,74,75,71,69,72,76,73]

Output:

[1,1,4,2,1,1,0,0]

Example 2

Input:

temps = [30,40,50,60]

Output:

[1,1,1,0]

Worked solution and explanation

What this really is

Strip the weather costume off and this is the classic next-greater-element scan: for each day you want the distance to the first later reading that beats it. Almost everyone reaches for the obvious double loop, look ahead from every day until you find something warmer, and it sails through the small examples. The trap is that it is O(n^2), and the interviewer is quietly waiting to see whether you notice that the days you already looked at carry information you keep throwing away. The fix is to hold the still-unresolved days on a stack of indices, warmest-waiting on top, and clear them in a single pass.

How to get there

Step 1: Seed the answer with zeros

Start every day at 0. That single default quietly handles the no-warmer-day case, so the loop only ever has to write the days that DO resolve.

Step 2: Resolve waits as the warmer day arrives

Walk the readings left to right, carrying a stack of indices whose temperatures strictly decrease from bottom to top. While today is warmer than the reading at the top index, that top day's wait just ended, so pop it and write today's index minus its index.

Step 3: Park the current day on the stack

Once the top is no longer colder than today, push today's index. It now waits its turn on the stack until some future warmer day pops it, or until the scan ends and it stays 0.

The solution

Monotonic stack, single pass
def daily_temperatures(temps: list) -> list:
    n = len(temps)
    result = [0] * n
    stack = []  # indices of days still waiting for a warmer one
    for i in range(n):
        while stack and temps[i] > temps[stack[-1]]:
            prev_idx = stack.pop()
            result[prev_idx] = i - prev_idx
        stack.append(i)
    return result

Each index is pushed once and popped at most once, so the whole thing is one linear sweep.

Brute force
From each day, scan forward until a warmer reading or the end. Correct, but a long cooling-then-warming run makes nearly every day rescan its successors: O(n^2). On a 10^5-day series this is the difference between instant and timing out.
Monotonic stack
Each day is examined a constant number of times total, because once a day is resolved it leaves the stack forever. The forward rescans collapse into a single left-to-right sweep.

Common follow-up questions