Source code for titan_pylib.algorithm.cycle_memo

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

[docs]
 6class CycleMemo(Generic[T]):
 7
 8    def __init__(self, initial: T, move_to: Callable[[T], T]) -> None:
 9        history: list[T] = []
10        memo: set[T] = set()
11        now: T = initial
12        while now not in memo:
13            history.append(now)
14            memo.add(now)
15            now = move_to(now)
16        cycle_start = history.index(now)
17        cycle_len = len(history) - cycle_start
18        pre = history[:cycle_start]
19        cycle = history[cycle_start:]
20        self.cycle_start = cycle_start
21        self.cycle_len = cycle_len
22        self.pre = pre
23        self.cycle = cycle
24

[docs]
25    def kth(self, k: int) -> T:
26        if k < self.cycle_start:
27            return self.pre[k]
28        k -= self.cycle_start
29        k %= self.cycle_len
30        return self.cycle[k]