1use std::{
2 borrow::Cow,
3 ops::{Index, IndexMut},
4};
5
6#[derive(Copy, Clone, Debug)]
7pub struct Arg(usize);
8
9impl From<()> for Arg {
10 fn from(_: ()) -> Self {
11 Self(32)
12 }
13}
14
15impl From<usize> for Arg {
16 fn from(initial_capacity: usize) -> Self {
17 Self(initial_capacity)
18 }
19}
20
21#[derive(Debug, Default)]
22pub struct TreeMap<K, V> {
23 nodes: Vec<Option<Node<K, V>>>,
24 hole_ids: Vec<usize>,
25 root: usize,
26}
27
28impl<K: Default + PartialOrd, V: Default, Iter: Iterator<Item = (K, V)>> From<Iter>
29 for TreeMap<K, V>
30{
31 fn from(iter: Iter) -> Self {
32 let mut this = Self::new(());
33 this.bulk_put(iter);
34 this
35 }
36}
37
38impl<K, V> TreeMap<K, V> {
39 pub fn new(arg: impl Into<Arg>) -> Self {
40 let Arg(initial_capacity) = arg.into();
41 let mut nodes = Vec::with_capacity(initial_capacity + 1);
42 nodes.push(None);
43 Self {
44 nodes,
45 hole_ids: Vec::with_capacity(initial_capacity),
46 root: 0,
47 }
48 }
49
50 pub fn is_empty(&self) -> bool {
51 self.len() == 0
52 }
53
54 pub fn len(&self) -> usize {
55 self.nodes.len() - 1 - self.hole_ids.len()
56 }
57
58 fn update(&mut self, id: usize) {
59 let mut cursor = id;
60 while cursor != 0 {
61 self.update_heights(cursor);
62 if self.is_biased(cursor) {
63 self.balance(cursor);
64 }
65 cursor = self.nodes[cursor].as_ref().unwrap().parent;
66 }
67 }
68
69 fn update_heights(&mut self, id: usize) {
70 self.update_left_height(id);
71 self.update_right_height(id);
72 }
73
74 fn update_left_height(&mut self, id: usize) {
75 let node = self.nodes[id].as_ref().unwrap();
76 let left = node.left;
77 self.nodes[id].as_mut().unwrap().left_height = if left == 0 {
78 0
79 } else {
80 let left_node = self.nodes[left].as_ref().unwrap();
81 left_node.left_height.max(left_node.right_height) + 1
82 }
83 }
84
85 fn update_right_height(&mut self, id: usize) {
86 let node = self.nodes[id].as_ref().unwrap();
87 let right = node.right;
88 self.nodes[id].as_mut().unwrap().right_height = if right == 0 {
89 0
90 } else {
91 let right_node = self.nodes[right].as_ref().unwrap();
92 right_node.left_height.max(right_node.right_height) + 1
93 }
94 }
95
96 fn is_biased(&self, id: usize) -> bool {
97 let node = self.nodes[id].as_ref().unwrap();
98 node.left_height.abs_diff(node.right_height) > 1
99 }
100
101 fn balance(&mut self, id: usize) {
102 let node = self.nodes[id].as_ref().unwrap();
103 if node.left_height < node.right_height {
104 let right_node = self.nodes[node.right].as_ref().unwrap();
105 if right_node.left_height < right_node.right_height {
106 self.rotate_left(id);
107 } else {
108 self.flip_right(id);
109 }
110 } else {
111 let left_node = self.nodes[node.left].as_ref().unwrap();
112 if left_node.left_height < left_node.right_height {
113 self.flip_left(id);
114 } else {
115 self.rotate_right(id);
116 }
117 }
118 }
119
120 fn rotate_left(&mut self, id: usize) {
121 let node = self.nodes[id].as_ref().unwrap();
122 let p = node.parent;
123 let r = node.right;
124 let rl = self.nodes[r].as_ref().unwrap().left;
125 if p == 0 {
126 self.root = r;
127 } else {
128 let p_node = self.nodes[p].as_mut().unwrap();
129 if p_node.left == id {
130 p_node.left = r;
131 } else {
132 p_node.right = r;
133 }
134 }
135 let node = self.nodes[id].as_mut().unwrap();
136 node.right = rl;
137 node.parent = r;
138 let r_node = self.nodes[r].as_mut().unwrap();
139 r_node.left = id;
140 r_node.parent = p;
141 if let Some(node) = self.nodes[rl].as_mut() {
142 node.parent = id;
143 }
144 self.update_right_height(id);
145 }
146
147 fn flip_right(&mut self, id: usize) {
148 let node = self.nodes[id].as_ref().unwrap();
149 let p = node.parent;
150 let r = node.right;
151 let rl = self.nodes[r].as_ref().unwrap().left;
152 let rl_node = self.nodes[rl].as_ref().unwrap();
153 let rll = rl_node.left;
154 let rlr = rl_node.right;
155 if p == 0 {
156 self.root = rl;
157 } else {
158 let p_node = self.nodes[p].as_mut().unwrap();
159 if p_node.left == id {
160 p_node.left = rl;
161 } else {
162 p_node.right = rl;
163 }
164 }
165 let r_node = self.nodes[r].as_mut().unwrap();
166 r_node.parent = rl;
167 r_node.left = rlr;
168 let rl_node = self.nodes[rl].as_mut().unwrap();
169 rl_node.parent = p;
170 rl_node.left = id;
171 rl_node.right = r;
172 let node = self.nodes[id].as_mut().unwrap();
173 node.parent = rl;
174 node.right = rll;
175 if let Some(node) = self.nodes[rll].as_mut() {
176 node.parent = id;
177 }
178 if let Some(node) = self.nodes[rlr].as_mut() {
179 node.parent = r;
180 }
181 self.update_right_height(id);
182 self.update_left_height(r);
183 }
184
185 fn rotate_right(&mut self, id: usize) {
186 let node = self.nodes[id].as_ref().unwrap();
187 let p = node.parent;
188 let l = node.left;
189 let lr = self.nodes[l].as_ref().unwrap().right;
190 if p == 0 {
191 self.root = l;
192 } else {
193 let p_node = self.nodes[p].as_mut().unwrap();
194 if p_node.left == id {
195 p_node.left = l;
196 } else {
197 p_node.right = l;
198 }
199 }
200 let node = self.nodes[id].as_mut().unwrap();
201 node.parent = l;
202 node.left = lr;
203 let l_node = self.nodes[l].as_mut().unwrap();
204 l_node.right = id;
205 l_node.parent = p;
206 if let Some(node) = self.nodes[lr].as_mut() {
207 node.parent = id;
208 }
209 self.update_left_height(id);
210 }
211
212 fn flip_left(&mut self, id: usize) {
213 let node = self.nodes[id].as_ref().unwrap();
214 let p = node.parent;
215 let l = node.left;
216 let lr = self.nodes[l].as_ref().unwrap().right;
217 let lr_node = self.nodes[lr].as_mut().unwrap();
218 let lrr = lr_node.right;
219 let lrl = lr_node.left;
220 let p_node = self.nodes[p].as_mut().unwrap();
221 if p == 0 {
222 self.root = lr;
223 } else if p_node.left == id {
224 p_node.left = lr;
225 } else {
226 p_node.right = lr;
227 }
228 let l_node = self.nodes[l].as_mut().unwrap();
229 l_node.parent = lr;
230 l_node.right = lrl;
231 let lr_node = self.nodes[lr].as_mut().unwrap();
232 lr_node.parent = p;
233 lr_node.right = id;
234 lr_node.left = l;
235 let node = self.nodes[id].as_mut().unwrap();
236 node.parent = lr;
237 node.left = lrr;
238 self.nodes[lrr].as_mut().unwrap().parent = id;
239 self.nodes[lrl].as_mut().unwrap().parent = l;
240 self.update_left_height(id);
241 self.update_right_height(l);
242 }
243
244 pub fn to_slice(&self) -> Cow<[(&K, &V)]> {
245 if self.root != 0 {
246 self.to_slice_from(self.root)
247 } else {
248 Cow::Borrowed(&[])
249 }
250 }
251
252 fn to_slice_from(&self, id: usize) -> Cow<[(&K, &V)]> {
253 let mut result = vec![];
254 let node = self.nodes[id].as_ref().unwrap();
255 if node.left != 0 {
256 result.extend_from_slice(&self.to_slice_from(node.left));
257 }
258 result.push((&node.key, &node.value));
259 if node.right != 0 {
260 result.extend_from_slice(&self.to_slice_from(node.right));
261 }
262 result.into()
263 }
264
265 pub fn clear(&mut self) {
266 self.nodes.clear();
267 self.nodes.push(None);
268 self.hole_ids.clear();
269 self.root = 0;
270 }
271}
272
273impl<K: Default, V: Default> TreeMap<K, V> {
274 fn insert_root(&mut self, key: K, value: V) {
275 if let Some(hole_id) = self.hole_ids.pop() {
276 self.fill_hole_with_root(hole_id, key, value);
277 } else {
278 self.push_root(key, value);
279 }
280 }
281
282 fn fill_hole_with_root(&mut self, hole_id: usize, key: K, value: V) {
283 self.nodes[hole_id] = Some(Node {
284 key,
285 value,
286 ..Default::default()
287 });
288 self.root = hole_id;
289 }
290
291 fn push_root(&mut self, key: K, value: V) {
292 self.nodes.push(Some(Node {
293 key,
294 value,
295 ..Default::default()
296 }));
297 self.root = 1;
298 }
299}
300
301impl<K: PartialOrd, V> TreeMap<K, V> {
302 pub fn has(&self, key: K) -> bool {
303 self.get(key).is_some()
304 }
305
306 pub fn get(&self, key: K) -> Option<&V> {
307 if let Some(id) = self.get_id(key) {
308 Some(&self.nodes[id].as_ref().unwrap().value)
309 } else {
310 None
311 }
312 }
313
314 pub fn get_mut(&mut self, key: K) -> Option<&mut V> {
315 if let Some(id) = self.get_id(key) {
316 Some(&mut self.nodes[id].as_mut().unwrap().value)
317 } else {
318 None
319 }
320 }
321
322 pub fn remove(&mut self, key: K) -> Option<V> {
323 if let Some(old) = self.get_id(key) {
324 let old_node_slot = self.nodes.get_mut(old).unwrap();
325 let old_node = old_node_slot.as_ref().unwrap();
326 let old_p = old_node.parent;
327 let old_l = old_node.left;
328 let old_r = old_node.right;
329 let result = old_node_slot.take().unwrap().value;
330 self.hole_ids.push(old);
331 if old_l != 0 {
332 let new = self.max_id_from(old_l);
333 let new_node = self.nodes[new].as_ref().unwrap();
334 let new_p = new_node.parent;
335 let new_l = new_node.left;
336 if old_p != 0 {
337 let old_p_node = self.nodes[old_p].as_mut().unwrap();
338 if old_p_node.left == old {
339 old_p_node.left = new;
340 } else {
341 old_p_node.right = new;
342 }
343 let new_node = self.nodes[new].as_mut().unwrap();
344 new_node.parent = old_p;
345 if new_p != old {
346 new_node.left = old_l;
347 }
348 new_node.right = old_r;
349 } else {
350 self.root = new;
351 let new_node = self.nodes[new].as_mut().unwrap();
352 new_node.parent = 0;
353 if new_p != old {
354 new_node.left = old_l;
355 }
356 new_node.right = old_r;
357 }
358 if old_r != 0 {
359 self.nodes[old_r].as_mut().unwrap().parent = new;
360 }
361 if new_p != old {
362 self.nodes[old_l].as_mut().unwrap().parent = new;
363 self.nodes[new_p].as_mut().unwrap().right = new_l;
364 if new_l != 0 {
365 self.nodes[new_l].as_mut().unwrap().parent = new_p;
366 }
367 self.update(new_p);
368 } else {
369 self.update(new);
370 }
371 } else if old_r != 0 {
372 let new = self.min_id_from(old_r);
373 let new_node = self.nodes[new].as_ref().unwrap();
374 let new_p = new_node.parent;
375 let new_r = new_node.right;
376 if old_p != 0 {
377 let old_p_node = self.nodes[old_p].as_mut().unwrap();
378 if old_p_node.left == old {
379 old_p_node.left = new;
380 } else {
381 old_p_node.right = new;
382 }
383 let new_node = self.nodes[new].as_mut().unwrap();
384 new_node.parent = old_p;
385 new_node.left = 0;
386 if new_p != old {
387 new_node.right = old_r;
388 }
389 } else {
390 self.root = new;
391 let new_node = self.nodes[new].as_mut().unwrap();
392 new_node.parent = 0;
393 new_node.left = 0;
394 if new_p != old {
395 new_node.right = old_r;
396 }
397 }
398 if new_p != old {
399 self.nodes[old_r].as_mut().unwrap().parent = new;
400 self.nodes[new_p].as_mut().unwrap().left = new_r;
401 if new_r != 0 {
402 self.nodes[new_r].as_mut().unwrap().parent = new_p;
403 }
404 self.update(new_p);
405 } else {
406 self.update(new);
407 }
408 } else if old_p != 0 {
409 let old_p_node = self.nodes[old_p].as_mut().unwrap();
410 if old_p_node.left == old {
411 old_p_node.left = 0;
412 } else {
413 old_p_node.right = 0;
414 }
415 self.update(old_p);
416 } else {
417 self.root = 0;
418 }
419 Some(result)
420 } else {
421 None
422 }
423 }
424
425 fn get_id(&self, key: K) -> Option<usize> {
426 let mut cursor = self.root;
427 while cursor != 0 {
428 let cursor_node = self.nodes[cursor].as_ref().unwrap();
429 if key < cursor_node.key {
430 cursor = cursor_node.left;
431 } else if key > cursor_node.key {
432 cursor = cursor_node.right;
433 } else {
434 return Some(cursor);
435 }
436 }
437 None
438 }
439
440 pub fn min(&self) -> Option<(&K, &V)> {
441 if self.root != 0 {
442 let min_node = self.nodes[self.min_id_from(self.root)].as_ref().unwrap();
443 Some((&min_node.key, &min_node.value))
444 } else {
445 None
446 }
447 }
448
449 fn min_id_from(&self, id: usize) -> usize {
450 let mut cursor = id;
451 loop {
452 let cursor_node = self.nodes[cursor].as_ref().unwrap();
453 if cursor_node.left != 0 {
454 cursor = cursor_node.left;
455 } else {
456 return cursor;
457 }
458 }
459 }
460
461 pub fn max(&self) -> Option<(&K, &V)> {
462 if self.root != 0 {
463 let max_node = self.nodes[self.max_id_from(self.root)].as_ref().unwrap();
464 Some((&max_node.key, &max_node.value))
465 } else {
466 None
467 }
468 }
469
470 fn max_id_from(&self, id: usize) -> usize {
471 let mut cursor = id;
472 loop {
473 let cursor_node = self.nodes[cursor].as_ref().unwrap();
474 if cursor_node.right != 0 {
475 cursor = cursor_node.right;
476 } else {
477 return cursor;
478 }
479 }
480 }
481}
482
483impl<K: PartialOrd, V> Index<K> for TreeMap<K, V> {
484 type Output = V;
485
486 fn index(&self, key: K) -> &Self::Output {
487 self
488 .get(key)
489 .expect("The value that is associated with the given key does not exist in the tree map!")
490 }
491}
492
493impl<K: PartialOrd, V> IndexMut<K> for TreeMap<K, V> {
494 fn index_mut(&mut self, key: K) -> &mut Self::Output {
495 self
496 .get_mut(key)
497 .expect("The value that is associated with the given key does not exist in the tree map!")
498 }
499}
500
501impl<K: Default + PartialOrd, V: Default> TreeMap<K, V> {
502 pub fn bulk_put(&mut self, iter: impl Iterator<Item = (K, V)>) {
503 for (key, value) in iter {
504 self.put(key, value);
505 }
506 }
507
508 pub fn put(&mut self, key: K, value: V) {
509 if self.root == 0 {
510 self.insert_root(key, value);
511 return;
512 }
513 let mut cursor = self.root;
514 loop {
515 let cursor_node = self.nodes[cursor].as_ref().unwrap();
516 if key < cursor_node.key {
517 if cursor_node.left == 0 {
518 let parent = cursor_node.parent;
519 self.insert_left(cursor, key, value);
520 self.update(parent);
521 break;
522 } else {
523 cursor = cursor_node.left;
524 }
525 } else if key > cursor_node.key {
526 if cursor_node.right == 0 {
527 let parent = cursor_node.parent;
528 self.insert_right(cursor, key, value);
529 self.update(parent);
530 break;
531 } else {
532 cursor = cursor_node.right;
533 }
534 } else {
535 self.nodes[cursor].as_mut().unwrap().value = value;
536 break;
537 }
538 }
539 }
540
541 fn insert_left(&mut self, id: usize, key: K, value: V) {
542 if let Some(hole_id) = self.hole_ids.pop() {
543 self.fill_hole_with_node_when_insert_left(id, hole_id, key, value);
544 } else {
545 self.push_node_when_insert_left(id, key, value);
546 }
547 }
548
549 fn fill_hole_with_node_when_insert_left(&mut self, id: usize, hole_id: usize, key: K, value: V) {
550 let node = self.nodes[id].as_mut().unwrap();
551 node.left = hole_id;
552 node.left_height += 1;
553 self.nodes[hole_id] = Some(Node {
554 key,
555 value,
556 parent: id,
557 ..Default::default()
558 });
559 }
560
561 fn push_node_when_insert_left(&mut self, id: usize, key: K, value: V) {
562 let node_count = self.nodes.len();
563 let node = self.nodes[id].as_mut().unwrap();
564 node.left = node_count;
565 node.left_height += 1;
566 self.nodes.push(Some(Node {
567 key,
568 value,
569 parent: id,
570 ..Default::default()
571 }));
572 }
573
574 fn insert_right(&mut self, id: usize, key: K, value: V) {
575 if let Some(hole_id) = self.hole_ids.pop() {
576 self.fill_hole_with_node_when_insert_right(id, hole_id, key, value);
577 } else {
578 self.push_node_when_insert_right(id, key, value);
579 }
580 }
581
582 fn fill_hole_with_node_when_insert_right(&mut self, id: usize, hole_id: usize, key: K, value: V) {
583 let node = self.nodes[id].as_mut().unwrap();
584 node.right = hole_id;
585 node.right_height += 1;
586 self.nodes[hole_id] = Some(Node {
587 key,
588 value,
589 parent: id,
590 ..Default::default()
591 });
592 }
593
594 fn push_node_when_insert_right(&mut self, id: usize, key: K, value: V) {
595 let node_count = self.nodes.len();
596 let node = self.nodes[id].as_mut().unwrap();
597 node.right = node_count;
598 node.right_height += 1;
599 self.nodes.push(Some(Node {
600 key,
601 value,
602 parent: id,
603 ..Default::default()
604 }));
605 }
606}
607
608#[derive(Debug, Default)]
609struct Node<K, V> {
610 key: K,
611 value: V,
612 parent: usize,
613 left: usize,
614 right: usize,
615 left_height: usize,
616 right_height: usize,
617}