Solution to LeetCode Next Greater Element Ⅰ in Python

This is my solution to LeetCode's "Next Greater Element Ⅰ":

You are given two arrays (without duplicates) nums1 and nums2 where nums1’s elements are subset of nums2. Find all the next greater numbers for nums1's elements in the corresponding places of nums2.

The Next Greater Number of a number x in nums1 is the first greater number to its right in nums2. If it does not exist, output -1 for this number.

class Solution:

def nextGreaterElement(self, nums1: List[int], nums2: List[int]) -> List[int]:
need, ready, next_greater = set(nums1), set(), {}
for n in nums2:
for k in ready:
if k < n:
next_greater[k] = n
ready = {k for k in ready if k not in next_greater}
if n in need and n not in next_greater and n not in ready:
return [next_greater[k] if k in next_greater else -1 for k in nums1]

The nextGreaterElement method accepts 2 lists. Neither list contains duplicates, and nums1 is a subset of nums2. For every num in nums1, it will output the first number greater than num positioned to the right of num in nums2, or -1 is none is found. e.g.:

Input: nums1 = [4,1,2], nums2 = [1,3,4,2].
Output: [-1,3,-1]

I loop through the keys in ready, and the body of the loop does nothing unless k < n. len(ready) can be very large relative to the number of values that satisfy k < n. This seems like a very common thing so I'm wondering if there's a more explicit/pythonic way to write this?

Or should I be using a different data structure entirely?

• @dfhwze My question is more about writing Python loops in general, not my specific implementation of the loop. Would StackOverflow be a better place for that kind of question? – NightDriveDrones Aug 2 at 20:19
• No this is the correct place for it. You added a description of the problem statement, so it's on-topic now. – dfhwze Aug 3 at 7:14

Your code seems unnecessarily complicated.

nums1 is subset of nums2. So no need to iterate over nums2. Because nums2 may be bigger than nums1

Here's the algo;

1. Take a number n from nums1
2. Find n's index in nums2. Because nums1 is subset of nums2 n will definitely be found.
3. Check if there's a number say m which is greater than n starting from the index + 1 index at nums2
4. If found output the number
5. Otherwise output -1

class Solution:
def nextGreaterElement(self, nums1: List[int], nums2: List[int]) -> List[int]:
output = []
for n in nums1:
idx = nums2.index(n)
for m in nums2[idx+1:]:
if m > n:
output.append(m)
break
else:
output.append(-1)
return output

Finding index sucks time. To optimize the index can be pre-calculated.

output = []
index = { n: i for i, n in enumerate(nums2) }
for n in nums1:
idx = index[n]
for m in nums2[idx+1:]:
if m > n:
output.append(m)
break
else:
output.append(-1)
return output

• Implementation details of Python's list.index might make this faster (I don't know), but the algorithm you're proposing is slower asymptotically as all elements of nums2 will be iterated for every element in nums1 (O(nm)) vs iterating each list once – NightDriveDrones Aug 2 at 21:56
• @user107870 that's good point – Nizam Mohamed Aug 2 at 22:04
• @user107870 reduced iteration – Nizam Mohamed Aug 3 at 0:37