Source code for titan_pylib.data_structures.stack.foldable_stack

 1from typing import Generic, Iterable, TypeVar, Callable, Union
 2
 3T = TypeVar("T")
 4
 5

[docs]
 6class FoldableStack(Generic[T]):
 7
 8    def __init__(
 9        self, n_or_a: Union[int, Iterable[T]], op: Callable[[T, T], T], e: T
10    ) -> None:
11        self._op = op
12        self._e = e
13        if isinstance(n_or_a, int):
14            self._n = n_or_a
15            self._a = [e] * self._n
16        else:
17            n_or_a = list(n_or_a)
18            self._n = len(n_or_a)
19            self._a = list(n_or_a)
20        self._data = [e] * (self._n + 1)
21        for i in range(self._n):
22            self._data[i + 1] = op(self._data[i], self._a[i])
23

[docs]
24    def append(self, key: T) -> None:
25        self._a.append(key)
26        self._data.append(self._op(self._data[-1], key))

27

[docs]
28    def top(self) -> T:
29        return self._a[-1]

30

[docs]
31    def pop(self) -> T:
32        self._data.pop()
33        return self._a.pop()

34

[docs]
35    def all_prod(self) -> T:
36        return self._data[-1]

37

[docs]
38    def prod(self, r: int) -> T:
39        return self._data[r]

40

[docs]
41    def tolist(self) -> list[T]:
42        return list(self._a)

43
44    def __len__(self):
45        return len(self._a)