[leetcode] 503. Next Greater Element II @ python

原题

Given a circular array (the next element of the last element is the first element of the array), print the Next Greater Number for every element. The Next Greater Number of a number x is the first greater number to its traversing-order next in the array, which means you could search circularly to find its next greater number. If it doesn’t exist, output -1 for this number.

Example 1:
Input: [1,2,1]
Output: [2,-1,2]
Explanation: The first 1’s next greater number is 2;
The number 2 can’t find next greater number;
The second 1’s next greater number needs to search circularly, which is also 2.

解法

堆栈. 由于题意要求循环查找更大值, 记录nums的原先长度prev_l, 我们将两个nums相连构成新的nums, 初始化ans为值为-1的列表. 遍历nums, 将暂未找到更大值的index, value放入堆栈, 如果当前n比堆栈里的值大, 那么一直出栈, 并更新ans. 如果没有找到更大值, 结果就是初始值-1, 返回ans的前prev_l的值.

Time: O(2n)
Space: O(2n)

代码

class Solution(object):
    def nextGreaterElements(self, nums):
        """
        :type nums: List[int]
        :rtype: List[int]
        """
        prev_l = len(nums)
        nums = nums+nums
        stack = []
        ans = [-1]*len(nums)
        for i, n in enumerate(nums):
            while stack and stack[-1][1] < n:
                prev_i = stack.pop()[0]
                ans[prev_i] = n
            stack.append([i, n])
            
        return ans[:prev_l]