Source code for titan_pylib.data_structures.stack.persistent_stack

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

[docs]
 6class PersistentStack(Generic[T]):
 7

[docs]
 8    class Node:
 9
10        def __init__(self, key: Optional[T], prev=None):
11            self.key: Optional[T] = key
12            self.prev: Optional[PersistentStack.Node] = prev

13
14    def __init__(
15        self,
16        a: Iterable[T] = [],
17        init_t: int = -1,
18        _prev: Optional["PersistentStack"] = None,
19    ):
20        self.prev: Optional[PersistentStack] = _prev
21        stack = PersistentStack.Node(None, None)
22        for e in a:
23            stack = PersistentStack.Node(e, prev=stack)
24        self.a: dict[int, PersistentStack.Node] = {init_t: stack}
25

[docs]
26    def top(self, t: int) -> T:
27        res = self.a[t].key
28        assert (
29            res is not None
30        ), f"IndexError: top() from empty {self.__class__.__name__}."
31        return res

32

[docs]
33    def append(self, key: T, pre_t: int, new_t: int) -> None:
34        s = self.a[pre_t]
35        new_stack = PersistentStack.Node(key, s)
36        self.a[new_t] = new_stack

37

[docs]
38    def pop(self, pre_t: int, new_t: int) -> T:
39        s = self.a[pre_t]
40        assert (
41            s.key is not None
42        ), f"IndexError: pop() from empty {self.__class__.__name__}."
43        self.a[new_t] = PersistentStack.Node(None) if s.prev is None else s.prev
44        return s.key

45

[docs]
46    def tolist(self, t: int) -> list[T]:
47        a = []
48        stack = self.a[t]
49        while stack.prev is not None:
50            a.append(stack.key)
51            stack = stack.prev
52        return a[::-1]