Update stock_span_problem.py

Author: MaximSmolskiy

data_structures/stacks/stock_span_problem.py

@@ -8,46 +8,71 @@
 """
 
 
-def calculation_span(price: list[float]) -> list[float]:
+def calculate_span(price: list[int]) -> list[int]:
     """
     Calculate the span values for a given list of stock prices.
     Args:
-        price (list): List of stock prices.
+        price: List of stock prices.
     Returns:
+        List of span values.
+
     >>> price = [10, 4, 5, 90, 120, 80]
-    >>> calculation_span(price)
+    >>> calculate_span(price)
     [1.0, 1.0, 2.0, 4.0, 5.0, 1.0]
     >>> price = [100, 50, 60, 70, 80, 90]
-    >>> calculation_span(price)
+    >>> calculate_span(price)
     [1.0, 1.0, 2.0, 3.0, 4.0, 5.0]
     >>> price = [5, 4, 3, 2, 1]
-    >>> calculation_span(price)
+    >>> calculate_span(price)
     [1.0, 1.0, 1.0, 1.0, 1.0]
     >>> price = [1, 2, 3, 4, 5]
-    >>> calculation_span(price)
+    >>> calculate_span(price)
     [1.0, 2.0, 3.0, 4.0, 5.0]
     >>> price = [10, 20, 30, 40, 50]
-    >>> calculation_span(price)
+    >>> calculate_span(price)
     [1.0, 2.0, 3.0, 4.0, 5.0]
     >>> calculation_span(price=[100, 80, 60, 70, 60, 75, 85])
     [1.0, 1.0, 1.0, 2.0, 1.0, 4.0, 6.0]
     """
     n = len(price)
-    st = [0]
-    s = [1.0]
+    s = [0] * n
+    # Create a stack and push index of fist element to it
+    st = []
+    st.append(0)
+
+    # Span value of first element is always 1
+    s[0] = 1
+
+    # Calculate span values for rest of the elements
     for i in range(1, n):
-        while st and price[st[-1]] <= price[i]:
+        # Pop elements from stack while stack is not
+        # empty and top of stack is smaller than price[i]
+        while len(st) > 0 and price[st[0]] <= price[i]:
             st.pop()
-        s.append(float(i - st[-1] if st else i + 1))
+
+        # If stack becomes empty, then price[i] is greater
+        # than all elements on left of it, i.e. price[0],
+        # price[1], ..price[i-1]. Else the price[i]  is
+        # greater than elements after top of stack
+        s[i] = i + 1 if len(st) <= 0 else (i - st[0])
+
+        # Push this element to stack
         st.append(i)
+
     return s
 
 
-price = [10.0, 4.0, 5.0, 90.0, 120.0, 80.0]
-S = calculation_span(price)
-print(S)
+# A utility function to print elements of array
+def print_array(arr, n):
+    for i in range(n):
+        print(arr[i], end=" ")
+
+
+# Driver program to test above function
+price = [10, 4, 5, 90, 120, 80]
 
-if __name__ == "__main__":
-    import doctest
+# Calculate the span values
+S = calculate_span(price)
 
-    doctest.testmod()
+# Print the calculated span values
+print_array(S, len(price))