Java 10 lines linear time complexity O(n) with explanation


  • 214
    Y

    Key observation:
    Suppose we have a decreasing sequence followed by a greater number
    For example [5, 4, 3, 2, 1, 6] then the greater number 6 is the next greater element for all previous numbers in the sequence

    We use a stack to keep a decreasing sub-sequence, whenever we see a number x greater than stack.peek() we pop all elements less than x and for all the popped ones, their next greater element is x
    For example [9, 8, 7, 3, 2, 1, 6]
    The stack will first contain [9, 8, 7, 3, 2, 1] and then we see 6 which is greater than 1 so we pop 1 2 3 whose next greater element should be 6

        public int[] nextGreaterElement(int[] findNums, int[] nums) {
            Map<Integer, Integer> map = new HashMap<>(); // map from x to next greater element of x
            Stack<Integer> stack = new Stack<>();
            for (int num : nums) {
                while (!stack.isEmpty() && stack.peek() < num)
                    map.put(stack.pop(), num);
                stack.push(num);
            }   
            for (int i = 0; i < findNums.length; i++)
                findNums[i] = map.getOrDefault(findNums[i], -1);
            return findNums;
        }
    

  • 0
    T

    Great solution!


  • 1

    Thanks. good solution!


  • 0
    F

    Great solution!


  • 7

    very smart solution. Here's a potential optimization that reduces your memory from O(nums length) to O(findNums length). The idea is you only need to hash elements that are in findNums all others can be ignored. This didn't show any improvement in OJ timer but certainly less memory is better.

        public int[] NextGreaterElement(int[] findNums, int[] nums) 
        {
            Dictionary<int,int> map = new Dictionary<int,int>();
            foreach (int x in findNums) map[x] = -1;
            
            Stack<int> stack = new Stack<int>();
            foreach (int x in nums)
            {
                while (stack.Count > 0 && stack.Peek() < x)
                {
                    map[stack.Pop()] = x;
                }
                if (map.ContainsKey(x)) stack.Push(x);
            }
            
            int[] res = new int[findNums.Length];
            for (int i = 0; i < findNums.Length; i++)
            {
                res[i] = map[findNums[i]];
            }
            
            return res;
        }
    

  • 2
    R

    Javascript version:

    var nextGreaterElement = function(findNums, nums) {
        var mp = {};
        var stack = []
        for(var i = 0; i < nums.length; i++) {
            
            while(stack.length > 0 && stack[stack.length-1] < nums[i])
                mp[stack.pop()] = nums[i];
            stack.push(nums[i])
        }
        for(var i = 0; i < findNums.length; i++) {
            if(findNums[i] in mp) {
                findNums[i] = mp[findNums[i]]    
            } else {
                findNums[i] = -1
            }       
        }
        return findNums;
    };
    

    Golang version:

    func nextGreaterElement(findNums []int, nums []int) []int {
        mp := make(map[int]int)
        stack := make([]int, 0)
        
        for i:= 0; i < len(nums); i++ {
            
            for len(stack) > 0 && stack[len(stack)-1] < nums[i] {
                last := stack[len(stack) - 1]
                stack = stack[:len(stack)-1]
                mp[last] = nums[i]
            }
            stack = append(stack, nums[i])
        }
        for i := 0; i < len(findNums); i++ {
            _, ok := mp[findNums[i]]
            if ok {
                findNums[i] = mp[findNums[i]]    
            } else {
                findNums[i] = -1
            }       
        }
        return findNums
    }
    

  • 1
    W

    Thanks for your post, and it is really a good solution.
    But it seems to me that it is not O(n). Could you explain why you think it is?
    Thanks in advance!


  • 2
    Y

    @weih1214
    Each element is pushed and popped at most once, that's why it's O(n)


  • 1
    A

    Man I am impressed, thanks for sharing such great solution!


  • 2

    Same idea. But I think loop from right to left is a little bit clearer.

    public class Solution {
        public int[] nextGreaterElement(int[] findNums, int[] nums) {
            Map<Integer, Integer> map = new HashMap<>();
            Stack<Integer> stack = new Stack<>();
            for(int i = nums.length-1; i>=0; i--){
                while(!stack.empty() && nums[i]>stack.peek()) stack.pop();
                map.put(nums[i], (stack.empty())? -1 : stack.peek());
                stack.push(nums[i]);
            }
            for(int i = 0; i<findNums.length; i++){
                findNums[i] = map.get(findNums[i]);
            }
            return findNums;        
        }
    }
    

  • 0
    S

    @yuxiangmusic Nice solution. This problem might be as well medium difficulty level.


  • 0

    Thanks, great solution !


  • 0
    M

    Really clever solution.
    This is the c++ version.

    class Solution {
    public:
        vector<int> nextGreaterElement(vector<int>& findNums, vector<int>& nums) {
            vector<int> ret;
            unordered_map<int, int> map_greater_element;
            stack<int> s;
            for (const auto& n : nums) {
                while (!s.empty() && s.top() < n) {
                    map_greater_element[s.top()] = n;
                    s.pop();
                }
                s.push(n);
            }
            for (size_t i = 0; i < findNums.size(); ++i) {
                unordered_map<int, int>::const_iterator cit = map_greater_element.find(findNums[i]);
                if (cit == map_greater_element.end()) {
                    ret.push_back(-1);
                } else {
                    ret.push_back(cit->second);
                }
            }
            return ret;
        }
    };
    

  • 0
    D
    This post is deleted!

  • 0
    A

    Python version (you could also avoid using the res list and instead change findNums in-place):

    class Solution(object):
        def nextGreaterElement(self, findNums, nums):
            """
            :type findNums: List[int]
            :type nums: List[int]
            :rtype: List[int]
            """
            stack = []
            d = {}
            res = []
            
            for num in nums:
                while stack and stack[-1] < num:
                    d[stack.pop(-1)] = num
                stack.append(num)
                
            for num in findNums:
                if num in d:
                    res.append(d[num])
                else:
                    res.append(-1)
            
            return res
    

  • 0
    R

    Clever! In the past 30 mins,I am thinking how to handle it by Stack...


  • 0
    C

    really osm... great


  • 0
    C

    Hi, why is my way of doing it without stack faster than this method? Thanks

    class Solution {
    public:
    vector<int> nextGreaterElement(vector<int>& findNums, vector<int>& nums) {
        unordered_map<int, int> numToIdx;
        
        for (int i=0; i<nums.size(); i++){
            numToIdx[nums[i]]=i;
        }
        
        vector<int> res;
        for (auto num : findNums){
            int idx = numToIdx[num];
            bool find = false;
            for (int i=idx+1; i<nums.size(); i++){
                if (nums[i]>num){
                    find=true;
                    res.push_back(nums[i]);
                    break;
                }
            }
            
            if (!find)
               res.push_back(-1);
        }
        
        return res;
    }
    

    };


  • 0
    T

    That is cool, bro!


  • 0
    H

    @jdrogin, Your algorithm's space Complexity is still O(nums length) in the worst case. You are using a stack to store elements from nums. In the worst case, it could be O(nums length).


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.