496. Next Greater Element I

Wrong Answer

class Solution:
    def nextGreaterElement(self, nums1: List[int], nums2: List[int]) -> List[int]:
        res = []

        for i in range(len(nums1)):
            idx = nums2.index(nums1[i])

            # error if last value in num2 == num1's some
            if nums2[-1] == nums2[idx]: # if match last value
                res.append(-1)
            elif nums2[idx] <= max(nums2[idx+1:]): # if next element is bigger
                res.append(max(nums2[idx+1:]))
            else:
                res.append(-1)

                # print(num2.index(num1[i]))
        return res

처음엔 이런식으로 풀려고 했는데 반례가 나오는 경우가 생김

Your input

[4,1,2]
[1,3,4,2]

Output

[-1,4,-1]

Expected

[-1,3,-1]

이런식으로 바로 옆의 값이 greater이면서 나머지 구간에서 max값이 아닌 경우 발생

 

class Solution:
    def nextGreaterElement(self, nums1: List[int], nums2: List[int]) -> List[int]:
        # nums1 is a subset of nums2
        result = []
        
        for i, n in list(enumerate(nums1)):
            ind = nums2.index(n)
            flag = False
            for i in range(ind, len(nums2)):
                if nums2[i] > n:
                    result.append(nums2[i])
                    flag = True
                    break
            if not flag:
                result.append('-1')
                
        return result

결국 for 문을 2번 돌려서 brute force 방식으로 푸는 방법이 있음

여기서 flag 사용해서 해야하는 것도 체크

 

Element 문제를 Elegant하게 푸는 방법

class Solution:
    def nextGreaterElement(self, nums1: List[int], nums2: List[int]) -> List[int]:
        
        stack = []
        d = {}
        
        for e in nums2:
            
            while(stack and e > stack[-1]): # if stack is not empty, and added value is greater than left
                c = stack.pop()             # pop 
                d[c] = e                    # insert dict as key=left value, value=right greater value
            
            stack.append(e)
        
        return [d.get(e, -1) for e in nums1] # get dict value from nums1, if not exist on dict then return -1