Source code for titan_pylib.data_structures.heap.double_ended_heap
1from titan_pylib.my_class.supports_less_than import SupportsLessThan
2from typing import Generic, Iterable, TypeVar
3
4T = TypeVar("T", bound=SupportsLessThan)
5
6
[docs]
7class DoubleEndedHeap(Generic[T]):
8 """
9 - [両端優先度付きキューのInterval-Heap実装](https://natsugiri.hatenablog.com/entry/2016/10/10/035445)
10 - [Double-ended priority queue(wikipedia)](https://en.wikipedia.org/wiki/Double-ended_priority_queue)
11 """
12
13 def __init__(self, a: Iterable[T] = []) -> None:
14 """構築します。
15 :math:`O(n)` です。
16
17 Args:
18 a (Iterable[T], optional): 構築元の配列です。
19 """
20 self._data = list(a)
21 self._heapify()
22
23 def _heapify(self) -> None:
24 n = len(self._data)
25 for i in range(n - 1, -1, -1):
26 if i & 1 and self._data[i - 1] < self._data[i]:
27 self._data[i - 1], self._data[i] = self._data[i], self._data[i - 1]
28 k = self._down(i)
29 self._up(k, i)
30
[docs]
31 def add(self, key: T) -> None:
32 """``key`` を1つ追加します。
33 :math:`O(\\log{n})` です。
34 """
35 self._data.append(key)
36 self._up(len(self._data) - 1)
37
[docs]
38 def pop_min(self) -> T:
39 """最小の要素を削除して返します。
40 :math:`O(\\log{n})` です。
41 """
42 if len(self._data) < 3:
43 res = self._data.pop()
44 else:
45 self._data[1], self._data[-1] = self._data[-1], self._data[1]
46 res = self._data.pop()
47 k = self._down(1)
48 self._up(k)
49 return res
50
[docs]
51 def pop_max(self) -> T:
52 """最大の要素を削除して返します。
53 :math:`O(\\log{n})` です。
54 """
55 if len(self._data) < 2:
56 res = self._data.pop()
57 else:
58 self._data[0], self._data[-1] = self._data[-1], self._data[0]
59 res = self._data.pop()
60 self._up(self._down(0))
61 return res
62
[docs]
63 def get_min(self) -> T:
64 """最小の要素を返します。
65 :math:`O(1)` です。
66 """
67 return self._data[0] if len(self._data) < 2 else self._data[1]
68
[docs]
69 def get_max(self) -> T:
70 """最大の要素を返します。
71 :math:`O(1)` です。
72 """
73 return self._data[0]
74
75 def __len__(self) -> int:
76 return len(self._data)
77
78 def __bool__(self) -> bool:
79 return len(self._data) > 0
80
81 def _parent(self, k: int) -> int:
82 return ((k >> 1) - 1) & ~1
83
84 def _down(self, k: int) -> int:
85 n = len(self._data)
86 if k & 1:
87 while k << 1 | 1 < n:
88 c = 2 * k + 3
89 if n <= c or self._data[c - 2] < self._data[c]:
90 c -= 2
91 if c < n and self._data[c] < self._data[k]:
92 self._data[k], self._data[c] = self._data[c], self._data[k]
93 k = c
94 else:
95 break
96 else:
97 while 2 * k + 2 < n:
98 c = 2 * k + 4
99 if n <= c or self._data[c] < self._data[c - 2]:
100 c -= 2
101 if c < n and self._data[k] < self._data[c]:
102 self._data[k], self._data[c] = self._data[c], self._data[k]
103 k = c
104 else:
105 break
106 return k
107
108 def _up(self, k: int, root: int = 1) -> int:
109 if (k | 1) < len(self._data) and self._data[k & ~1] < self._data[k | 1]:
110 self._data[k & ~1], self._data[k | 1] = (
111 self._data[k | 1],
112 self._data[k & ~1],
113 )
114 k ^= 1
115 while root < k:
116 p = self._parent(k)
117 if not self._data[p] < self._data[k]:
118 break
119 self._data[p], self._data[k] = self._data[k], self._data[p]
120 k = p
121 while root < k:
122 p = self._parent(k) | 1
123 if not self._data[k] < self._data[p]:
124 break
125 self._data[p], self._data[k] = self._data[k], self._data[p]
126 k = p
127 return k
128
[docs]
129 def tolist(self) -> list[T]:
130 return sorted(self._data)
131
132 def __str__(self) -> str:
133 return str(self.tolist())
134
135 def __repr__(self) -> str:
136 return f"{self.__class__.__name__}({self})"