diff --git a/data_structures/stacks/next_greater_element.py b/data_structures/stacks/next_greater_element.py index 216850b4b894..62d14d3a14cb 100644 --- a/data_structures/stacks/next_greater_element.py +++ b/data_structures/stacks/next_greater_element.py @@ -1,133 +1,28 @@ -from __future__ import annotations - -arr = [-10, -5, 0, 5, 5.1, 11, 13, 21, 3, 4, -21, -10, -5, -1, 0] -expect = [-5, 0, 5, 5.1, 11, 13, 21, -1, 4, -1, -10, -5, -1, 0, -1] - - -def next_greatest_element_slow(arr: list[float]) -> list[float]: - """ - Get the Next Greatest Element (NGE) for each element in the array - by checking all subsequent elements to find the next greater one. - - This is a brute-force implementation, and it has a time complexity - of O(n^2), where n is the size of the array. - - Args: - arr: List of numbers for which the NGE is calculated. - - Returns: - List containing the next greatest elements. If no - greater element is found, -1 is placed in the result. - - Example: - >>> next_greatest_element_slow(arr) == expect - True - """ - - result = [] - arr_size = len(arr) - - for i in range(arr_size): - next_element: float = -1 - for j in range(i + 1, arr_size): - if arr[i] < arr[j]: - next_element = arr[j] - break - result.append(next_element) - return result - - -def next_greatest_element_fast(arr: list[float]) -> list[float]: - """ - Find the Next Greatest Element (NGE) for each element in the array - using a more readable approach. This implementation utilizes - enumerate() for the outer loop and slicing for the inner loop. - - While this improves readability over next_greatest_element_slow(), - it still has a time complexity of O(n^2). - - Args: - arr: List of numbers for which the NGE is calculated. - - Returns: - List containing the next greatest elements. If no - greater element is found, -1 is placed in the result. - - Example: - >>> next_greatest_element_fast(arr) == expect - True - """ - result = [] - for i, outer in enumerate(arr): - next_item: float = -1 - for inner in arr[i + 1 :]: - if outer < inner: - next_item = inner - break - result.append(next_item) - return result - - -def next_greatest_element(arr: list[float]) -> list[float]: - """ - Efficient solution to find the Next Greatest Element (NGE) for all elements - using a stack. The time complexity is reduced to O(n), making it suitable - for larger arrays. - - The stack keeps track of elements for which the next greater element hasn't - been found yet. By iterating through the array in reverse (from the last - element to the first), the stack is used to efficiently determine the next - greatest element for each element. - - Args: - arr: List of numbers for which the NGE is calculated. - - Returns: - List containing the next greatest elements. If no - greater element is found, -1 is placed in the result. - - Example: - >>> next_greatest_element(arr) == expect - True - """ - arr_size = len(arr) - stack: list[float] = [] - result: list[float] = [-1] * arr_size - - for index in reversed(range(arr_size)): +""" +Next Greater Element +-------------------- +Given an array, find the next greater element for each element in the array. +The next greater element for an element x is the first greater element on the +right side of x in the array. If no such element exists, output -1. +""" + + +def next_greater_elements(arr): + stack = [] + result = [-1] * len(arr) + + # Traverse from right to left + for i in range(len(arr) - 1, -1, -1): + while stack and stack[-1] <= arr[i]: + stack.pop() if stack: - while stack[-1] <= arr[index]: - stack.pop() - if not stack: - break - if stack: - result[index] = stack[-1] - stack.append(arr[index]) + result[i] = stack[-1] + stack.append(arr[i]) return result if __name__ == "__main__": - from doctest import testmod - from timeit import timeit - - testmod() - print(next_greatest_element_slow(arr)) - print(next_greatest_element_fast(arr)) - print(next_greatest_element(arr)) - - setup = ( - "from __main__ import arr, next_greatest_element_slow, " - "next_greatest_element_fast, next_greatest_element" - ) - print( - "next_greatest_element_slow():", - timeit("next_greatest_element_slow(arr)", setup=setup), - ) - print( - "next_greatest_element_fast():", - timeit("next_greatest_element_fast(arr)", setup=setup), - ) - print( - " next_greatest_element():", - timeit("next_greatest_element(arr)", setup=setup), - ) + sample = [4, 5, 2, 25] + print("Input:", sample) + print("Next Greater Elements:", next_greater_elements(sample)) + # Output: [5, 25, 25, -1]