Adding doctests in stock_span_problem.py

data_structures/stacks/stock_span_problem.py

@@ -9,43 +9,62 @@
 
 
 def calculation_span(price, s):
+    """
+    Calculate the span values for a given list of stock prices.
+    Args:
+        price (list): List of stock prices.
+        s (list): List to store the span values.
+    Returns:
+        None
+    >>> price = [10, 4, 5, 90, 120, 80]
+    >>> S = [0 for i in range(len(price) + 1)]
+    >>> calculation_span(price, S)
+    >>> S
+    [1, 1, 2, 4, 5, 6, 0]
+    >>> price = [100, 50, 60, 70, 80, 90]
+    >>> S = [0 for i in range(len(price) + 1)]
+    >>> calculation_span(price, S)
+    >>> S
+    [1, 1, 2, 3, 4, 5, 0]
+    >>> price = [5, 4, 3, 2, 1]
+    >>> S = [0 for i in range(len(price) + 1)]
+    >>> calculation_span(price, S)
+    >>> S
+    [1, 1, 2, 3, 4, 0]
+    >>> price = [1, 2, 3, 4, 5]
+    >>> S = [0 for i in range(len(price) + 1)]
+    >>> calculation_span(price, S)
+    >>> S
+    [1, 2, 3, 4, 5, 0]
+    >>> price = [10, 20, 30, 40, 50]
+    >>> S = [0 for i in range(len(price) + 1)]
+    >>> calculation_span(price, S)
+    >>> S
+    [1, 2, 3, 4, 5, 0]
+    """
     n = len(price)
-    # 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):
-        # 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()
-
-        # 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)
 
-
-# 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]
 S = [0 for i in range(len(price) + 1)]
 
-# Fill the span values in array S[]
 calculation_span(price, S)
 
-# Print the calculated span values
 print_array(S, len(price))
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()