Source code for titan_pylib.data_structures.bst_base.bst_multiset_node_base
1from typing import TypeVar, Generic, Optional
2
3T = TypeVar("T")
4Node = TypeVar("Node")
5# protcolで、key,val,left,right を規定
6
7
[docs]
8class BSTMultisetNodeBase(Generic[T, Node]):
9
[docs]
10 @staticmethod
11 def count(node, key: T) -> int:
12 while node is not None:
13 if node.key == key:
14 return node.val
15 node = node.left if key < node.key else node.right
16 return 0
17
[docs]
18 @staticmethod
19 def get_min(node: Node) -> Optional[T]:
20 if not node:
21 return None
22 while node.left:
23 node = node.left
24 return node.key
25
[docs]
26 @staticmethod
27 def get_max(node: Node) -> Optional[T]:
28 if not node:
29 return None
30 while node.right:
31 node = node.right
32 return node.key
33
[docs]
34 @staticmethod
35 def contains(node: Node, key: T) -> bool:
36 while node:
37 if key == node.key:
38 return True
39 node = node.left if key < node.key else node.right
40 return False
41
[docs]
42 @staticmethod
43 def tolist(node: Node) -> list[T]:
44 stack = []
45 a = []
46 while stack or node:
47 if node:
48 stack.append(node)
49 node = node.left
50 else:
51 node = stack.pop()
52 for _ in range(node.val):
53 a.append(node.key)
54 node = node.right
55 return a
56
[docs]
57 @staticmethod
58 def tolist_items(node: Node) -> list[tuple[T, int]]:
59 stack = []
60 a = []
61 while stack or node:
62 if node:
63 stack.append(node)
64 node = node.left
65 else:
66 node = stack.pop()
67 a.append((node.key, node.val))
68 node = node.right
69 return a
70
71 @staticmethod
72 def _rle(a: list[T]) -> tuple[list[T], list[int]]:
73 keys, vals = [], []
74 keys.append(a[0])
75 vals.append(1)
76 for i, elm in enumerate(a):
77 if i == 0:
78 continue
79 if elm == keys[-1]:
80 vals[-1] += 1
81 continue
82 keys.append(elm)
83 vals.append(1)
84 return keys, vals
85
[docs]
86 @staticmethod
87 def le(node: Node, key: T) -> Optional[T]:
88 res = None
89 while node is not None:
90 if key == node.key:
91 res = key
92 break
93 if key < node.key:
94 node = node.left
95 else:
96 res = node.key
97 node = node.right
98 return res
99
[docs]
100 @staticmethod
101 def lt(node: Node, key: T) -> Optional[T]:
102 res = None
103 while node is not None:
104 if key <= node.key:
105 node = node.left
106 else:
107 res = node.key
108 node = node.right
109 return res
110
[docs]
111 @staticmethod
112 def ge(node: Node, key: T) -> Optional[T]:
113 res = None
114 while node is not None:
115 if key == node.key:
116 res = key
117 break
118 if key < node.key:
119 res = node.key
120 node = node.left
121 else:
122 node = node.right
123 return res
124
[docs]
125 @staticmethod
126 def gt(node: Node, key: T) -> Optional[T]:
127 res = None
128 while node is not None:
129 if key < node.key:
130 res = node.key
131 node = node.left
132 else:
133 node = node.right
134 return res
135
[docs]
136 @staticmethod
137 def index(node: Node, key: T) -> int:
138 k = 0
139 while node:
140 if key == node.key:
141 if node.left:
142 k += node.left.valsize
143 break
144 if key < node.key:
145 node = node.left
146 else:
147 k += node.val if node.left is None else node.left.valsize + node.val
148 node = node.right
149 return k
150
[docs]
151 @staticmethod
152 def index_right(node: Node, key: T) -> int:
153 k = 0
154 while node:
155 if key == node.key:
156 k += node.val if node.left is None else node.left.valsize + node.val
157 break
158 if key < node.key:
159 node = node.left
160 else:
161 k += node.val if node.left is None else node.left.valsize + node.val
162 node = node.right
163 return k