Next Greater Element I | Solution Explained + Stack Optimization

LeetCode 496. Next Greater Element I is a simple problem, since brute-force solution can get accepted. But, we are also going to look at a stack based optimization. Problem Link.

0. Next Greater Element I: Brute-force Solution (Accepted)

We do exactly as the question says.

• Iterate over all the elements in nums1, say the current is nums1[i], the ith element
• Then, find the j in nums2 that matches; nums1[i] == nums2[j]
• Then, find the k in nums2 from [j, n) that is nums2[k] > nums1[i]
• Add it to the answer.

class Solution:
    def nextGreaterElement(self, nums1: List[int], nums2: List[int]) -> List[int]:
        m, n = len(nums1), len(nums2)
        
        ans = []
        for i in range(m):
            for j in range(n):
                if nums1[i] == nums2[j]: 
                    for k in range(j, n):
                        if nums2[k] > nums1[i]:
                            ans.append(nums2[k])
                            break
                if len(ans) == i+1: # if already found for jth index in nums2
                    break
            if len(ans) != i+1: # if not found
                ans.append(-1)
        
        return ans

1. Next Greater Element I: Optimization Logic

1.0. Observations

To find the answers for each element of nums1, we want to know the answer for each element of nums2. So, first let’s try to find that.

These are some samples of what can happen in nums2.

1.1. Black Box

The only data structure that fits all the characteristics that we want is the stack.

1.2. Stack dry run

Explained in detail in the video.

2. Next Greater Element I: Optimized Code

class Solution:
    def nextGreaterElement(self, nums1: List[int], nums2: List[int]) -> List[int]:
        d = defaultdict(int)
        stack = []
        
        for num in nums2:  # O(M)
            while len(stack) and num > stack[-1]: 
                d[stack.pop()] = num
            stack.append(num)
        
        while stack:
            d[stack.pop()] = -1
        
        # O(N)
        return [d[num] for num in nums1]

2.0. Complexity Analysis

• Time: O(M+N). We have 2 operations for each of the num in nums2:
  • adding the element to the stack
  • removing the element from the stack
  • thus, O(2*N) time for creating the dictionary
• Space: O(N) The size of the dictionary is O(N).

Tanishq Chaudhary

Producing high-quality intuitive explanations of interview problems. Currently covering LeetCode and InterviewBit.