Leetcode: 2454. Next Greater Element IV

Problem Statement

Understand the problem

Given an array of non-negative integers, find for each integer in the array the second greatest element to its right.

Derive a plan

Let's keep two sorted containers \(a\) and \(b\). Elements on \(a\) didn't see a greater element to the left, and elements on \(b\) saw one element greater than them to their left. For each element \(i\) in the array (left to right), we remove from \(b\) all elements that are smaller than \(i\) since \(i\) is the second greatest element to their left; after that we transfer elements from \(a\) to \(b\) that are smaller than \(i\). Finally, we add \(i\) to \(a\) and continue the process. Time complexity is \(O(n \log n)\) and space is \(O(n)\).

Carry out the plan

from sortedcontainers import SortedList


class Solution:
    def secondGreaterElement(self, nums: List[int]) -> List[int]:
        N = len(nums)
        a = SortedList([], key=lambda i: nums[i])
        b = SortedList([], key=lambda i: nums[i])
        ans = [-1] * N
        a.add(0)
        for i in range(1, N):
            while b and nums[b[0]] < nums[i]:
                ans[b.pop(0)] = nums[i]

            while a and nums[a[0]] < nums[i]:
                b.add(a.pop(0))

            a.add(i)

        return ans