1use core::fmt;
2use std::{fmt::Debug, ops::Index, vec::IntoIter};
3
4pub struct SortedList<T>
35where
36 T: Ord,
37{
38 _lists: Vec<Vec<T>>,
39 _index_tree: Vec<usize>,
40 _index_tree_offset: usize,
41 _load_factor: usize,
42 _upper_load_factor: usize,
43 _lower_load_factor: usize,
44 _len: usize,
45}
46
47impl<T> SortedList<T>
49where
50 T: Ord,
51{
52 const DEFAULT_INDEX_TREE_OFFSET: usize = 1 << 5;
53 const DEFAULT_LOAD_FACTOR: usize = 1_024;
54 const DEFAULT_UPPER_LOAD_FACTOR: usize = 2_048;
55 const DEFAULT_LOWER_LOAD_FACTOR: usize = 512;
56
57 fn _default() -> Self {
59 Self {
60 _lists: vec![],
61 _index_tree: vec![0; 2 * Self::DEFAULT_INDEX_TREE_OFFSET],
62 _index_tree_offset: Self::DEFAULT_INDEX_TREE_OFFSET,
63 _load_factor: Self::DEFAULT_LOAD_FACTOR,
64 _upper_load_factor: Self::DEFAULT_UPPER_LOAD_FACTOR,
65 _lower_load_factor: Self::DEFAULT_LOWER_LOAD_FACTOR,
66 _len: 0,
67 }
68 }
69
70 fn _collapse(&mut self, i: usize) {
72 if self._lists.len() <= 1 {
73 panic!(
74 "Attempting to collapse while self._lists contains only {} lists.",
75 self._lists.len()
76 );
77 }
78
79 let left = match i >= 1 {
80 true => self._lists[i - 1].len(),
81 false => usize::MAX,
82 };
83
84 let right = match i + 1 < self._lists.len() {
85 true => self._lists[i + 1].len(),
86 false => usize::MAX,
87 };
88
89 assert!(left.min(right) < usize::MAX);
90 if left < right {
91 let mut removed = self._lists.remove(i);
93 self._lists[i - 1].append(&mut removed);
94 } else {
95 let mut removed = self._lists.remove(i + 1);
96 self._lists[i].append(&mut removed);
97 }
98
99 self._rebuild_index_tree();
100 }
101
102 fn _expand(&mut self, i: usize) {
104 if self._lists[i].len() < self._upper_load_factor {
105 panic!("Unnecessary expand at self._lists[{}]", i);
106 }
107
108 let size = self._lists[i].len();
109 let removed: Vec<T> = self._lists[i].drain(size / 2..).collect();
110 self._lists.insert(i + 1, removed);
111
112 self._rebuild_index_tree();
116 }
117
118 fn _rebuild_index_tree(&mut self) {
120 self._index_tree_offset = Self::DEFAULT_INDEX_TREE_OFFSET; while self._index_tree_offset < self._lists.len() {
122 self._index_tree_offset *= 2;
123 }
124
125 self._index_tree.fill(0);
126 self._index_tree.resize(2 * self._index_tree_offset, 0);
127
128 (0..self._lists.len()).for_each(|node| {
129 self._index_tree[node + self._index_tree_offset] = self._lists[node].len();
130 });
131
132 (1..self._index_tree_offset).rev().for_each(|node| {
133 self._index_tree[node] = self._index_tree[2 * node] + self._index_tree[2 * node + 1];
134 });
135 }
136
137 fn _index_tree_sum(
140 &self,
141 ql: usize,
142 qr: usize,
143 opt_node: Option<usize>,
144 opt_l: Option<usize>,
145 opt_r: Option<usize>,
146 ) -> usize {
147 let node = opt_node.unwrap_or(1);
148 let l = opt_l.unwrap_or(0);
149 let r = opt_r.unwrap_or(self._index_tree_offset - 1);
150
151 if ql <= l && r <= qr {
152 return self._index_tree[node];
153 }
154
155 if qr < l || r < ql {
156 return 0;
157 }
158
159 let m = (l + r) / 2;
160 self._index_tree_sum(ql, qr, Some(2 * node), Some(l), Some(m))
161 + self._index_tree_sum(ql, qr, Some(2 * node + 1), Some(m + 1), Some(r))
162 }
163
164 fn _index_tree_add(&mut self, i: usize, val: i32) {
166 let mut node = self._index_tree_offset + i;
167 if val >= 0 {
168 self._index_tree[node] += val as usize;
169 } else {
170 self._index_tree[node] -= (-val) as usize;
171 }
172 node /= 2;
173
174 while node > 0 {
175 self._index_tree[node] = self._index_tree[2 * node] + self._index_tree[2 * node + 1];
176 node /= 2;
177 }
178 }
179
180 fn _lists_remove(&mut self, i: usize, j: usize) -> T {
182 if i >= self._lists.len() || j >= self._lists[i].len() {
183 panic!(
184 "List index out of range. Attempting to remove self._lists[{}][{}]",
185 i, j
186 );
187 }
188
189 let removed = self._lists[i].remove(j);
190 self._len -= 1;
191
192 if self._lists.len() > 1 && self._lists[i].len() < self._lower_load_factor {
193 self._collapse(i);
194 } else {
195 self._index_tree_add(i, -1);
196 }
197
198 removed
199 }
200
201 fn _lists_insert(&mut self, i: usize, element: T) {
203 let pos = match self._lists[i].binary_search(&element) {
208 Ok(p) => p,
209 Err(p) => p,
210 };
211
212 self._lists[i].insert(pos, element);
213 self._len += 1;
214
215 if self._lists[i].len() > self._upper_load_factor {
216 self._expand(i);
217 } else {
218 self._index_tree_add(i, 1);
219 }
220 }
221
222 fn _bisect_right_lists(&self, element: &T) -> usize {
224 if &self._lists[0][0] > element {
225 return 0;
226 }
227
228 let mut lo = 0;
229 let mut hi = self._lists.len() - 1;
230 if &self._lists[hi][0] <= element {
231 return hi;
232 }
233
234 let mut mid;
237 while lo + 1 < hi {
238 mid = (lo + hi) / 2;
239 if &self._lists[mid][0] <= element {
240 lo = mid;
241 } else {
242 hi = mid;
243 }
244 }
245
246 lo
247 }
248
249 fn _locate_kth_element(&self, k: usize) -> (usize, usize) {
251 if k >= self._len {
253 panic!("SortedList: Index out of range.");
254 }
255
256 let is_leaf_node = |u| u >= self._index_tree_offset;
257 let mut cnt = k + 1;
258
259 let mut node: usize = 1;
260 while !is_leaf_node(node) {
261 if self._index_tree[2 * node] >= cnt {
262 node *= 2;
263 } else {
264 cnt -= self._index_tree[2 * node];
265 node = 2 * node + 1;
266 }
267 }
268
269 (node - self._index_tree_offset, cnt - 1)
271 }
272
273 fn _at(&self, i: usize, j: usize) -> &T {
275 &self._lists[i][j]
276 }
277
278 fn _flat(&self) -> Vec<&T> {
280 self._lists.iter().fold(Vec::new(), |mut cur, list| {
281 list.iter().for_each(|element| {
282 cur.push(element);
283 });
284 cur
285 })
286 }
287}
288
289impl<T> SortedList<T>
291where
292 T: Ord,
293{
294 pub fn new() -> Self {
304 Self::_default()
305 }
306
307 pub fn kth_smallest(&self, k: usize) -> &T {
318 let (i, j) = self._locate_kth_element(k);
320 return self._at(i, j);
321 }
322
323 pub fn clear(&mut self) {
337 self._lists.clear();
338 self._index_tree.clear();
339 self._index_tree
340 .resize(2 * Self::DEFAULT_INDEX_TREE_OFFSET, 0);
341 self._index_tree_offset = Self::DEFAULT_INDEX_TREE_OFFSET;
342 self._len = 0;
343 }
344
345 pub fn insert(&mut self, element: T) {
362 if self._len == 0 {
363 self._lists.clear();
364 self._lists.push(vec![]);
365 self._lists_insert(0, element);
366 return;
367 }
368
369 let k = self._bisect_right_lists(&element);
370 self._lists_insert(k, element);
371 }
372
373 pub fn remove(&mut self, k: usize) -> T {
386 let (i, j) = self._locate_kth_element(k);
387 self._lists_remove(i, j)
388 }
389
390 pub fn binary_search(&self, element: &T) -> Result<usize, usize> {
407 if self._len == 0 {
408 return Err(0);
409 }
410
411 let i: usize = self._bisect_right_lists(element);
412 if i == 0 {
413 return self._lists[i].binary_search(element);
414 }
415
416 match self._lists[i].binary_search(element) {
417 Ok(pos) => Ok(pos + self._index_tree_sum(0, i - 1, None, None, None)),
418 Err(pos) => Err(pos + self._index_tree_sum(0, i - 1, None, None, None)),
419 }
420 }
421
422 pub fn contains(&self, element: &T) -> bool {
435 self.binary_search(element).is_ok()
436 }
437
438 pub fn len(&self) -> usize {
450 self._len
451 }
452
453 pub fn is_empty(&self) -> bool {
467 self.len() == 0
468 }
469
470 pub fn last(&self) -> Option<&T> {
482 if self.is_empty() {
483 return None;
484 }
485 return self._lists.last().unwrap().last();
486 }
487
488 pub fn first(&self) -> Option<&T> {
500 if self.is_empty() {
501 return None;
502 }
503 return self._lists.first().unwrap().first();
504 }
505
506 pub fn get(&self, index: usize) -> Option<&T> {
518 if self.is_empty() || self.len() <= index {
519 return None;
520 }
521 return Some(self.kth_smallest(index));
522 }
523
524 pub fn flatten(&self) -> Vec<&T> {
537 self._flat()
538 }
539
540 pub fn to_vec(&self) -> Vec<T>
553 where
554 T: Clone,
555 {
556 self._lists.iter().fold(vec![], |mut cur, list| {
557 cur.append(&mut list.to_vec());
558 cur
559 })
560 }
561
562 pub fn iter(&self) -> impl Iterator<Item = &T> {
579 self._lists.iter().flat_map(|list| list.iter())
580 }
581}
582
583impl<T> Default for SortedList<T>
584where
585 T: Ord,
586{
587 fn default() -> Self {
589 Self::_default()
590 }
591}
592
593impl<T> Index<usize> for SortedList<T>
594where
595 T: Ord,
596{
597 type Output = T;
598
599 fn index(&self, index: usize) -> &Self::Output {
611 self.kth_smallest(index)
612 }
613}
614
615impl<T> From<IntoIter<T>> for SortedList<T>
616where
617 T: Ord,
618{
619 fn from(iter: IntoIter<T>) -> Self {
621 let mut array: Vec<T> = iter.collect();
622 array.sort();
623
624 let sorted_iter = array.into_iter();
627 let mut sorted_list = Self::default();
628 sorted_list._len = sorted_iter.len();
629
630 sorted_list._lists.push(vec![]);
631 let mut last_list_size = 0;
632
633 for element in sorted_iter {
634 sorted_list._lists.last_mut().unwrap().push(element);
635 last_list_size += 1;
636 if last_list_size == sorted_list._load_factor {
637 last_list_size = 0;
638 sorted_list._lists.push(vec![]);
639 }
640 }
641
642 sorted_list._rebuild_index_tree();
643 sorted_list
644 }
645}
646
647impl<T> From<Vec<T>> for SortedList<T>
648where
649 T: Ord,
650{
651 fn from(array: Vec<T>) -> Self {
653 Self::from(array.into_iter())
654 }
655}
656
657impl<T> From<&[T]> for SortedList<T>
658where
659 T: Ord + Clone,
660{
661 fn from(array: &[T]) -> Self {
663 Self::from(Vec::from(array))
664 }
665}
666
667impl<T> From<&mut [T]> for SortedList<T>
668where
669 T: Ord + Clone,
670{
671 fn from(array: &mut [T]) -> Self {
673 Self::from(Vec::from(array))
674 }
675}
676
677impl<T, const N: usize> From<[T; N]> for SortedList<T>
678where
679 T: Ord,
680{
681 fn from(array: [T; N]) -> Self {
683 Self::from(Vec::from(array))
684 }
685}
686
687impl<T> fmt::Debug for SortedList<T>
688where
689 T: Ord + Debug,
690{
691 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
692 fmt::Debug::fmt(&self._flat(), f)
693 }
694}
695
696#[cfg(test)]
697mod tests {
698 use rand::{seq::SliceRandom, thread_rng, Rng};
699
700 use crate::SortedList;
701
702 #[test]
703 fn it_works() {
704 let x = vec![1, 2, 3, 4, 9, 8, 7, 6, 5, 4];
705 let y = vec![1, 2, 3, 4, 4, 5, 6, 7, 8, 9];
706 let arr = SortedList::from(x);
707
708 for i in 0..10 {
709 assert_eq!(y[i], arr[i]);
710 }
711 }
712
713 #[test]
714 fn random_tests() {
715 let test_size = 100_000;
718 let test_op = 5_000;
719
720 let mut rng = thread_rng();
721 let mut array = vec![];
722 for _ in 0..test_size {
723 let x = rng.gen::<i32>();
724 array.push(x);
725 }
726
727 let mut copy = array.clone();
729 copy.sort();
730
731 let mut sorted_list = SortedList::from(array);
733
734 for _ in 0..test_op {
736 let x = rng.gen::<i32>();
737
738 let k = match copy.binary_search(&x) {
740 Ok(i) => i,
741 Err(i) => i,
742 };
743 copy.insert(k, x);
744
745 sorted_list.insert(x);
747 }
748
749 for _ in 0..test_op + test_size {
751 let idx = rng.gen_range(0..copy.len());
753 let expect = copy.remove(idx);
754 let actual = sorted_list.remove(idx);
755 assert_eq!(expect, actual);
756
757 let expect = copy.first();
759 let actual = sorted_list.first();
760 assert_eq!(expect, actual);
761
762 let expect = copy.last();
764 let actual = sorted_list.last();
765 assert_eq!(expect, actual);
766
767 let x = rng.gen::<i32>();
769 let actual = sorted_list.binary_search(&x);
770 let expect = copy.binary_search(&x);
771 assert_eq!(expect, actual);
772
773 let index = rng.gen_range(0..copy.len() + 2000);
775 let actual = sorted_list.get(index);
776 let expect = copy.get(index);
777 assert_eq!(expect, actual);
778
779 let actual = sorted_list.len();
781 let expect = copy.len();
782 assert_eq!(expect, actual);
783
784 let actual = sorted_list.is_empty();
786 let expect = copy.is_empty();
787 assert_eq!(expect, actual);
788 }
789 }
790
791 #[test]
792 fn example() {
793 let array = vec![90, 19, 25];
794 let mut sorted_list = SortedList::from(array);
795
796 sorted_list.insert(100);
797 sorted_list.insert(1);
798 sorted_list.insert(20);
799
800 let x = sorted_list.remove(3);
801 assert_eq!(25, x);
802 assert_eq!(&20, sorted_list.kth_smallest(2));
803 assert_eq!(20, sorted_list[2]);
804 }
805
806 #[test]
807 fn binary_search_test() {
808 let array = vec![20; 100_000];
809 let sorted_list = SortedList::from(array);
810 let x = 50;
811
812 let actual = sorted_list.binary_search(&x);
813 let expected: Result<usize, usize> = Err(100_000);
814
815 assert_eq!(actual, expected);
816 }
817
818 #[test]
819 fn contains_test() {
820 let mut sorted_list = SortedList::from([10; 10_000]);
821 assert!(sorted_list.contains(&10));
822 assert!(!sorted_list.contains(&9));
823
824 sorted_list.insert(9);
825 assert!(sorted_list.contains(&9));
826
827 sorted_list.insert(11);
828 assert!(sorted_list.contains(&11));
829 }
830
831 #[test]
832 fn clear_test() {
833 let mut sorted_list_1 = SortedList::from([10; 10_000]);
835 let mut sorted_list_2 = SortedList::from([1, 3, 4, 2, 0]);
836
837 sorted_list_1.clear();
839 sorted_list_2.clear();
840
841 for sorted_list in [sorted_list_1, sorted_list_2] {
843 assert!(sorted_list._lists.is_empty());
844 for val in &sorted_list._index_tree {
845 assert!(val == &0);
846 }
847 assert!(sorted_list._index_tree.len() == 2 * sorted_list._index_tree_offset);
848 }
849 }
850
851 #[test]
852 fn to_vec_test() {
853 let mut rng = thread_rng();
855 let mut array: Vec<usize> = (0..5_000).collect();
856 array.shuffle(&mut rng);
857
858 let sorted_list = SortedList::from(array);
859
860 let to_vec = sorted_list.to_vec();
862
863 assert_eq!((0..5_000).collect::<Vec<usize>>(), to_vec);
865 }
866
867 #[test]
868 fn flatten_test() {
869 let mut rng = thread_rng();
871 let mut array: Vec<usize> = (0..5_000).collect();
872 array.shuffle(&mut rng);
873
874 let sorted_list = SortedList::from(array);
875
876 let flatten = sorted_list.flatten();
878
879 for i in 0..5_000 {
881 assert_eq!(flatten[i], &i);
882 }
883 }
884
885 #[test]
886 fn complex_data_structure_test() {
887 let mut array: Vec<Vec<i32>> = vec![];
889 for i in (0..2000).rev() {
890 array.push((0..i + 1).collect());
891 }
892
893 let sorted_list = SortedList::from(array);
895
896 for i in 0..2000 {
898 let actual: &Vec<i32> = &sorted_list[i];
899 let expected = (0..i + 1).map(|x| x as i32).collect::<Vec<i32>>();
900
901 assert_eq!(i + 1, actual.len());
902 for j in 0..actual.len() {
903 assert_eq!(actual[j], expected[j]);
904 }
905 }
906 }
907
908 #[test]
909 fn break_case_insert_after_lst_has_been_clean() {
910 let mut lst = SortedList::<usize>::new();
911 lst.insert(3);
912 lst.remove(0);
913 lst.insert(1);
914 lst.insert(5);
915 }
916}