1# from titan_pylib.data_structures.heap.deletable_max_heap import DeletableMaxHeap
2# from titan_pylib.data_structures.heap.max_heap import MaxHeap
3# from titan_pylib.my_class.supports_less_than import SupportsLessThan
4from typing import Protocol
5
6
7class SupportsLessThan(Protocol):
8
9 def __lt__(self, other) -> bool: ...
10from typing import TypeVar, Generic, Iterable
11
12T = TypeVar("T", bound=SupportsLessThan)
13
14
15class MaxHeap(Generic[T]):
16
17 def __init__(self, a: Iterable[T] = []) -> None:
18 self.a = list(a)
19 self._heapify()
20
21 def _heapify(self) -> None:
22 for i in range(len(self.a) - 1, -1, -1):
23 self._down(i)
24
25 def _down(self, i: int) -> None:
26 a = self.a
27 n = len(a)
28 while i * 2 + 1 < n:
29 u, v = i * 2 + 1, i * 2 + 2
30 if v < n and a[u] < a[v]:
31 u = v
32 if a[i] < a[u]:
33 a[i], a[u] = a[u], a[i]
34 i = u
35 else:
36 break
37
38 def _up(self, i: int) -> None:
39 a = self.a
40 while i > 0:
41 p = (i - 1) >> 1
42 if a[i] > a[p]:
43 a[i], a[p] = a[p], a[i]
44 i = p
45 else:
46 break
47
48 def get_max(self) -> T:
49 return self.a[0]
50
51 def pop_max(self) -> T:
52 res = self.a[0]
53 self.a[0] = self.a[-1]
54 self.a.pop()
55 self._down(0)
56 return res
57
58 def push(self, key: T) -> None:
59 self.a.append(key)
60 self._up(len(self.a) - 1)
61
62 def pushpop_max(self, key: T) -> T:
63 if self.a[0] < key or self.a[0] == key:
64 return key
65 res = self.a[0]
66 self.a[0] = key
67 self._down(0)
68 return res
69
70 def replace_max(self, key: T) -> T:
71 res = self.a[0]
72 self.a[0] = key
73 self._down(0)
74 return res
75
76 def tolist(self) -> list[T]:
77 return sorted(self.a)
78
79 def __getitem__(self, k: int) -> T:
80 assert k == 0
81 return self.a[0]
82
83 def __len__(self):
84 return len(self.a)
85
86 def __str__(self):
87 return str(self.a)
88
89 def __repr__(self):
90 return f"{self.__class__.__name__}({self})"
91# from titan_pylib.my_class.supports_less_than import SupportsLessThan
92from typing import TypeVar, Generic, Iterable
93
94T = TypeVar("T", bound=SupportsLessThan)
95
96
97class DeletableMaxHeap(Generic[T]):
98
99 def __init__(self, a: Iterable[T] = []) -> None:
100 """削除可能Maxヒープです。
101 要素の 追加/削除/最大値取得 が効率よく行えます。
102 """
103 self.hq: MaxHeap[T] = MaxHeap(a)
104 self.rem_hq: MaxHeap[T] = MaxHeap()
105 self._len: int = len(self.hq)
106
107 def push(self, key: T) -> None:
108 """``key`` を追加します。
109 :math:`O(\\log{n})` です。
110 """
111 self._len += 1
112 if self.rem_hq and self.rem_hq.get_max() == key:
113 self.rem_hq.pop_max()
114 return
115 self.hq.push(key)
116
117 def remove(self, key: T) -> None:
118 """``key`` を削除します。
119 :math:`O(\\log{n})` です。
120 """
121 assert self._len > 0
122 self._len -= 1
123 if self.hq.get_max() == key:
124 self.hq.pop_max()
125 else:
126 self.rem_hq.push(key)
127
128 def _delete(self) -> None:
129 while self.rem_hq and self.rem_hq.get_max() == self.hq.get_max():
130 self.hq.pop_max()
131 self.rem_hq.pop_max()
132
133 def get_max(self) -> T:
134 """最大の要素を返します。
135 :math:`O(\\log{n})` です。
136 """
137 assert self._len > 0
138 self._delete()
139 return self.hq.get_max()
140
141 def pop_max(self) -> T:
142 """最大の要素を削除して返します。
143 :math:`O(\\log{n})` です。
144 """
145 assert self._len > 0
146 self._len -= 1
147 self._delete()
148 return self.hq.pop_max()
149
150 def __len__(self) -> int:
151 return self._len