Sorted List
A treap-based sorted list data structure.
Supports: - O(log n) insert, remove, bisect_left, bisect_right, get - O(1) len, is_empty
Uses a flat array representation (indices instead of boxed nodes) with a free stack for memory reuse.
Rust Implementation
Generated API reference for the Sorted List crate. View the rendered API reference →
1//! A treap-based sorted list.
2//!
3//! Supports the following operations, each `O(log n)`:
4//! - [`SortedList::insert`] / [`SortedList::remove`]
5//! - [`SortedList::bisect_left`] / [`SortedList::bisect_right`]
6//! - [`SortedList::get`]
7//!
8//! [`SortedList::len`] and [`SortedList::is_empty`] are `O(1)`.
9//!
10//! Implementation note: nodes are stored in a flat `Vec<Node>` addressed by
11//! 1-based indices (0 = none). Removed nodes go to a `free_stack` so the
12//! underlying storage can be reused without growing the array.
13
14use std::cmp::Ordering;
15
16#[derive(Debug)]
17struct Node {
18 key: i64,
19 priority: u32,
20 left: usize,
21 right: usize,
22 size: usize,
23}
24
25impl Node {
26 fn new(key: i64, seed: &mut u32) -> Self {
27 *seed ^= *seed << 13;
28 *seed ^= *seed >> 17;
29 *seed ^= *seed << 5;
30 Self {
31 key,
32 priority: *seed,
33 left: 0,
34 right: 0,
35 size: 1,
36 }
37 }
38
39 fn update(idx: usize, nodes: &mut [Node]) {
40 if idx == 0 {
41 return;
42 }
43 let l = nodes[idx - 1].left;
44 let r = nodes[idx - 1].right;
45 let l_size = if l == 0 { 0 } else { nodes[l - 1].size };
46 let r_size = if r == 0 { 0 } else { nodes[r - 1].size };
47 nodes[idx - 1].size = 1 + l_size + r_size;
48 }
49}
50
51/// Treap-backed sorted list of `i64` keys with multiplicity.
52pub struct SortedList {
53 nodes: Vec<Node>,
54 root: usize,
55 seed: u32,
56 free_stack: Vec<usize>,
57}
58
59impl SortedList {
60 /// Construct an empty sorted list.
61 pub fn new() -> Self {
62 Self {
63 nodes: Vec::with_capacity(1024 * 8),
64 root: 0,
65 seed: 12345,
66 free_stack: Vec::new(),
67 }
68 }
69
70 /// Split by key. If `inclusive` is true, keys == key go to the left tree.
71 /// Returns `(left_idx, right_idx)` where indices are 1-based (0 means None).
72 fn split(&mut self, node_idx: usize, key: i64, inclusive: bool) -> (usize, usize) {
73 if node_idx == 0 {
74 return (0, 0);
75 }
76
77 let n_key = self.nodes[node_idx - 1].key;
78 let go_left = if inclusive { n_key >= key } else { n_key > key };
79
80 if go_left {
81 let left_child = self.nodes[node_idx - 1].left;
82 let (l, r) = self.split(left_child, key, inclusive);
83 self.nodes[node_idx - 1].left = r;
84 Node::update(node_idx, &mut self.nodes);
85 (l, node_idx)
86 } else {
87 let right_child = self.nodes[node_idx - 1].right;
88 let (l, r) = self.split(right_child, key, inclusive);
89 self.nodes[node_idx - 1].right = l;
90 Node::update(node_idx, &mut self.nodes);
91 (node_idx, r)
92 }
93 }
94
95 /// Merge two trees where all keys in `l` are less than all keys in `r`.
96 fn merge(&mut self, l: usize, r: usize) -> usize {
97 if l == 0 || r == 0 {
98 return if l == 0 { r } else { l };
99 }
100
101 if self.nodes[l - 1].priority > self.nodes[r - 1].priority {
102 let l_right = self.nodes[l - 1].right;
103 self.nodes[l - 1].right = self.merge(l_right, r);
104 Node::update(l, &mut self.nodes);
105 l
106 } else {
107 let r_left = self.nodes[r - 1].left;
108 self.nodes[r - 1].left = self.merge(l, r_left);
109 Node::update(r, &mut self.nodes);
110 r
111 }
112 }
113
114 /// Insert a key. Duplicates are allowed.
115 pub fn insert(&mut self, key: i64) {
116 let (l, r) = self.split(self.root, key, true);
117 let node_idx = if let Some(idx) = self.free_stack.pop() {
118 self.nodes[idx - 1] = Node::new(key, &mut self.seed);
119 idx
120 } else {
121 self.nodes.push(Node::new(key, &mut self.seed));
122 self.nodes.len()
123 };
124 let merged_left = self.merge(l, node_idx);
125 self.root = self.merge(merged_left, r);
126 }
127
128 /// Remove one occurrence of `key`.
129 pub fn remove(&mut self, key: i64) {
130 let (l, mid_r) = self.split(self.root, key, true);
131 let (mid, r) = self.split(mid_r, key + 1, true);
132 if mid != 0 {
133 let m_left = self.nodes[mid - 1].left;
134 let m_right = self.nodes[mid - 1].right;
135 let new_mid = self.merge(m_left, m_right);
136 self.free_stack.push(mid);
137 let merged_left = self.merge(l, new_mid);
138 self.root = self.merge(merged_left, r);
139 } else {
140 self.root = self.merge(l, r);
141 }
142 }
143
144 /// Number of keys strictly less than `key`.
145 pub fn bisect_left(&self, key: i64) -> usize {
146 let mut curr = self.root;
147 let mut rank = 0;
148 while curr != 0 {
149 let node = &self.nodes[curr - 1];
150 if node.key >= key {
151 curr = node.left;
152 } else {
153 let l_size = if node.left == 0 { 0 } else { self.nodes[node.left - 1].size };
154 rank += 1 + l_size;
155 curr = node.right;
156 }
157 }
158 rank
159 }
160
161 /// Number of keys less than or equal to `key`.
162 pub fn bisect_right(&self, key: i64) -> usize {
163 let mut curr = self.root;
164 let mut rank = 0;
165 while curr != 0 {
166 let node = &self.nodes[curr - 1];
167 if node.key > key {
168 curr = node.left;
169 } else {
170 let l_size = if node.left == 0 { 0 } else { self.nodes[node.left - 1].size };
171 rank += 1 + l_size;
172 curr = node.right;
173 }
174 }
175 rank
176 }
177
178 /// Get the element at `index` (0-based), or `None` if out of bounds.
179 pub fn get(&self, index: usize) -> Option<i64> {
180 let mut curr = self.root;
181 let mut i = index;
182 while curr != 0 {
183 let node = &self.nodes[curr - 1];
184 let left_size = if node.left == 0 { 0 } else { self.nodes[node.left - 1].size };
185 match i.cmp(&left_size) {
186 Ordering::Less => curr = node.left,
187 Ordering::Equal => return Some(node.key),
188 Ordering::Greater => {
189 i -= left_size + 1;
190 curr = node.right;
191 }
192 }
193 }
194 None
195 }
196
197 /// Number of elements in the list (with multiplicities).
198 pub fn len(&self) -> usize {
199 if self.root == 0 {
200 0
201 } else {
202 self.nodes[self.root - 1].size
203 }
204 }
205
206 /// Whether the list is empty.
207 pub fn is_empty(&self) -> bool {
208 self.len() == 0
209 }
210}