Repository: TheAlgorithms/Python
File: data_structures/stacks/stock_span_problem.py (Python)

Code before:
number of consecutive days just before the given day, for which the price of the stock
on the current day is less than or equal to its price on the given day.
"""


def calculation_span(price: list[float]) -> list[float]:
    """
    Calculate the span values for a given list of stock prices.
    Args:
        price (list): List of stock prices.
    Returns:
    >>> price = [10, 4, 5, 90, 120, 80]
    >>> calculation_span(price)
    [1.0, 1.0, 2.0, 4.0, 5.0, 6.0]
    >>> price = [100, 50, 60, 70, 80, 90]
    >>> calculation_span(price)
    [1.0, 1.0, 2.0, 3.0, 4.0, 5.0]
    >>> price = [5, 4, 3, 2, 1]
    >>> calculation_span(price)
    [1.0, 1.0, 2.0, 3.0, 4.0]
    >>> price = [1, 2, 3, 4, 5]
    >>> calculation_span(price)
    [1.0, 2.0, 3.0, 4.0, 5.0]
    >>> price = [10, 20, 30, 40, 50]
    >>> calculation_span(price)
    [1.0, 2.0, 3.0, 4.0, 5.0]
    """
    n = len(price)
    st = [0]
    s = [0.0] * n
    s[0] = 1.0
    for i in range(1, n):
        while len(st) > 0 and price[st[0]] <= price[i]:
            st.pop()
        s[i] = float(i + 1) if len(st) <= 0 else float(i - st[0])
    return s


price = [10.0, 4.0, 5.0, 90.0, 120.0, 80.0]
S = calculation_span(price)
print(S)

if __name__ == "__main__":
    import doctest

    doctest.testmod()


Review comment (suggestion):
The output does not agree with the results in https://www.geeksforgeeks.org/the-stock-span-problem/
```suggestion
    [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]
    >>> calculation_span(price=[10, 4, 5, 90, 120, 80])
    [1.0, 1.0, 2.0, 4.0, 5.0, 1.0]
```

Code after:
number of consecutive days just before the given day, for which the price of the stock
on the current day is less than or equal to its price on the given day.
"""


def calculate_span(price: list[int]) -> list[int]:
    """
    Calculate the span values for a given list of stock prices.
    Args:
        price: List of stock prices.
    Returns:
        List of span values.

    >>> calculate_span([10, 4, 5, 90, 120, 80])
    [1, 1, 2, 4, 5, 1]
    >>> calculate_span([100, 50, 60, 70, 80, 90])
    [1, 1, 2, 3, 4, 5]
    >>> calculate_span([5, 4, 3, 2, 1])
    [1, 1, 1, 1, 1]
    >>> calculate_span([1, 2, 3, 4, 5])
    [1, 2, 3, 4, 5]
    >>> calculate_span([10, 20, 30, 40, 50])
    [1, 2, 3, 4, 5]
    >>> calculate_span([100, 80, 60, 70, 60, 75, 85])
    [1, 1, 1, 2, 1, 4, 6]
    """
    n = len(price)
    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):
        # 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[-1]] <= 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[-1])

        # Push this element to stack
        st.append(i)