Problem statement

https://leetcode.com/problems/next-greater-element-i/

Solution

Sizes of both arrays are small enough, so we just can do brute-force solution in O(m * n), where n is size of nums2 and m is size of nums1.

If we want to solve this problem in O(n) time, it is not so simple. The idea is to traverse nums2 and keep stack with decreasing order of elements. When we try to add element, if it is less than last element of stack, we just add it. If it is more than the last element, we extract it from stack and also put it inside dic: correspondence between numbers and its next greater element: we need it, because we have also nums1, which we need to traverse after. Next, when we traverse nums1 we can use function .get(num, -1), which will return answer for num if it is inside dictionary and -1 if it was not found.

Complexity

Time and space complexity is O(n).

Code

class Solution:
    def nextGreaterElement(self, nums1, nums2):
        dic, stack = {}, []
        
        for num in nums2[::-1]:
            while stack and num > stack[-1]:
                stack.pop()
            if stack:
                dic[num] = stack[-1]
            stack.append(num)
            
        return [dic.get(num, -1) for num in nums1]