[Swift 알고리즘] LeetCode 496. Next Greater Element I

○

Stack, HashTable | Easy

496. Next Greater Element I

문제
The next greater element of some element x in an array is the first greater element that is to the right of x in the same array.
You are given two distinct 0-indexed integer arrays nums1 and nums2, where nums1 is a subset of nums2.
For each 0 <= i < nums1.length, find the index j such that nums1[i] == nums2[j] and determine the next greater element of nums2[j] in nums2. If there is no next greater element, then the answer for this query is -1.
Return an array ans of length nums1.length such that ans[i] is the next greater element as described above.
Constraints
- 1 <= nums1.length <= nums2.length <= 1000
- 0 <= nums1[i], nums2[i] <= 10^4
- All integers in nums1 and nums2 are unique.
- All the integers of nums1 also appear in nums2.

출력 :

print(solution.nextGreaterElement([4,1,2], [1,3,4,2]))//[-1,3,-1]
print(solution.nextGreaterElement([2,4], [2,1,3,4]))//[3,-1]
print(solution.nextGreaterElement([4,3,6,2,5], [5,2,1,6,3,4]))//[-1,4,-1,6,6]

 

1. O(n^2)

브루트포스로 그냥 O(n^2)으로 푸는 건 쉽다. 통과도 가능하긴 하다.

func nextGreaterElement(_ nums1: [Int], _ nums2: [Int]) -> [Int] {
    var ans: [Int] = []
    
    for num in nums1 {
        ans.append(-1)
        let index = nums2.firstIndex(of: num)! 
        for i in (index+1)..<nums2.count {
            if num < nums2[i] {
                ans[ans.count - 1] = nums2[i]
                break
            }
        }
    }
    
    return ans
}

2. O(n)

그치만 출제의도(?) 에 맞게 스택을 사용해서 O(n+m)으로 풀려면 좀 생각이 필요한 문제다.

처음에 nums1 의 순서와 디폴트 값 -1을 먼저 딕셔너리에 넣고 nums2 를 탐색하려니 굉장히 복잡했는데,

nums2 에서의 next greater element 를 먼저 다 구해서 딕셔너리에 넣은 후 num1에서 값을 넣어주니 훨씬 깔끔하더라.

nums2 먼저 탐색하면서 해시에 저장하는 것이 관건인 문제인 듯. 

class Solution {
    func nextGreaterElement(_ nums1: [Int], _ nums2: [Int]) -> [Int] {
        var dict: [Int:Int] = [:]
        var stack: [Int] = []
        var ans: [Int] = []
        
        for num in nums2 {
            while let last = stack.last, num > last {
                dict[last] = num
                stack.removeLast()
            }
            stack.append(num)
        }
        
        for num in nums1 {
            ans.append(dict[num] ?? -1)
        }
        
        return ans
    }
}

처음에 nums2 탐색해서 딕셔너리 넣어주는 부분을 아래와 같이 구현했었는데, 

위 처럼 while let 을 사용해서 만들어주니 훨씬 간단하더라. while let 처음 써봄. 이런 것도 있구나..! 좋구만기레

for num in nums2 {
    for _ in 0..<stack.count {
        if let last = stack.last {
            if num > last {
                dict[last] = num
                stack.removeLast()
            } else { break }
        }
    }
    stack.append(num)
}