739. Daily Temperatures: Monotonic Stack

I initially couldn’t think of a better way than a O(nlogn) binary search method (O(nlogn) a list that sorts by (temp, day index) and O(logn) lookup per day). Peeked on the optimal data stucture to use which was a (monotonic) stack. Time to understand this.

Simply put, a monotonic stack is a stack but you maintain it so that it’s kept sorted either asc or desc. This is surprisingly useful on “retrieve the next greater/smaller element” questions.

We want to design our stack according to the question like this: we want to quickly locate the next day that is hotter than today. This induces us to construct the stack from the newer days to the older days since we need to know the future.

bottom [newest, newer, new, ...] top

We want the nearest day as long as it’s hotter, and a stack already solved that question for us, since we LIFO and push from the newer days, the newest day will always be at the bottom.

I asked for a hint after my first code failed to incorporate the date alongside the temps (how are you supposed to calculate the date distance in O(1) per day). We need to keep track of the (temp, idx).

 1class Solution:
 2    def dailyTemperatures(self, temperatures: List[int]) -> List[int]:
 3        n = len(temperatures)
 4        result = []
 5        stack = []  # (temp, date_idx)
 6
 7        for i in range(n - 1, -1, -1):
 8            temp = temperatures[i]
 9
10            while stack and stack[-1][0] <= temp:
11                stack.pop()
12
13            if stack:
14                result.append(stack[-1][1] - i)
15            else:
16                result.append(0)
17            stack.append((temp, i))
18
19        return result[::-1]

This is optimal O(n) TC SC. We can still make small improvements:

 1class Solution:
 2    def dailyTemperatures(self, temperatures: List[int]) -> List[int]:
 3        n = len(temperatures)
 4        result = [0] * n    # prevented the need to flip the result at the end
 5        stack = []  # we only store idx since temperatures[i] == temp
 6
 7        for i in range(n - 1, -1, -1):
 8            while stack and temperatures[stack[-1]] <= temperatures[i]:
 9                stack.pop()
10
11            result[i] = stack[-1] - i if stack else 0
12            stack.append(i)
13
14        return result