1use crate::{merge_entry::MergeEntry, merge_tree::MergeTree};
2
3pub struct LoserTree<T>
4where
5 T: Ord,
6{
7 losers: Vec<usize>,
9 leaves: Vec<Option<MergeEntry<T>>>,
11 winner: usize,
13 size: usize,
14}
15
16impl<T: Clone + Ord + PartialOrd> LoserTree<T> {
17 pub fn with_capacity(capacity: usize) -> Self {
18 let size = capacity.next_power_of_two();
19 let mut losers = Vec::with_capacity(size - 1);
20 let mut leaves = Vec::with_capacity(size);
21
22 losers.resize(size - 1, 0);
23 leaves.resize(size, None);
24 Self {
25 losers,
26 leaves,
27 winner: 0,
28 size: 0,
29 }
30 }
31
32 #[inline]
34 fn play_match(&mut self, pos1: usize, pos2: usize) -> usize {
35 use std::cmp::Ordering;
36
37 match (&self.leaves[pos1], &self.leaves[pos2]) {
38 (None, _) => pos2,
39 (_, None) => pos1,
40 (Some(v1), Some(v2)) => match v1.cmp(v2) {
41 Ordering::Greater => pos1,
42 _ => pos2,
43 },
44 }
45 }
46
47 fn rebuild_path(&mut self, mut i: usize) {
49 let n = self.leaves.len();
50 let internal_nodes = n - 1;
51
52 i += internal_nodes;
54
55 while i > 0 {
57 let parent = (i - 1) / 2;
59
60 let sibling = if i % 2 == 0 { i - 1 } else { i + 1 };
62
63 let winner = self.play_match(
65 if i >= internal_nodes {
66 i - internal_nodes
67 } else {
68 self.losers[i]
69 },
70 if sibling >= internal_nodes {
71 sibling - internal_nodes
72 } else {
73 self.losers[sibling]
74 },
75 );
76
77 self.losers[parent] = winner;
78 i = parent;
79 }
80
81 self.winner = self.losers[0];
82 }
83
84 pub fn clear(&mut self) {
85 self.size = 0;
86 self.winner = 0;
87 }
88
89 pub fn capacity(&self) -> usize {
90 self.leaves.len()
91 }
92}
93
94impl<T: Ord + std::clone::Clone> MergeTree<T> for LoserTree<T> {
95 fn push(&mut self, item: MergeEntry<T>) {
96 debug_assert!(item.index < self.leaves.len(), "index out of bounds");
97
98 let index = item.index;
99 self.leaves[index] = Some(item);
100
101 self.rebuild_path(index);
102 self.size += 1;
103 }
104
105 fn pop(&mut self) -> std::option::Option<MergeEntry<T>> {
106 let winner = self.leaves[self.winner].take();
107 if winner.is_some() {
108 self.size -= 1;
109 self.rebuild_path(self.winner); }
111 winner
112 }
113
114 fn is_empty(&self) -> bool {
115 self.size == 0
116 }
117
118 fn len(&self) -> usize {
119 self.size
120 }
121}
122
123#[cfg(test)]
124mod tests {
125 use super::*;
126
127 #[test]
128 fn test_new_loser_tree() {
129 let tree: LoserTree<i32> = LoserTree::with_capacity(3);
130 assert_eq!(tree.leaves.len(), 4); assert_eq!(tree.losers.len(), 3); assert!(tree.is_empty());
133 assert_eq!(tree.len(), 0);
134 }
135
136 #[test]
137 fn test_push_and_pop() {
138 let mut tree = LoserTree::with_capacity(4);
139
140 tree.push(MergeEntry { item: 3, index: 0 });
141 tree.push(MergeEntry { item: 1, index: 1 });
142 tree.push(MergeEntry { item: 4, index: 2 });
143 tree.push(MergeEntry { item: 2, index: 3 });
144
145 assert_eq!(tree.len(), 4);
146 assert!(!tree.is_empty());
147
148 assert_eq!(tree.pop().unwrap().item, 1);
150 assert_eq!(tree.pop().unwrap().item, 2);
151 assert_eq!(tree.pop().unwrap().item, 3);
152 assert_eq!(tree.pop().unwrap().item, 4);
153
154 assert!(tree.is_empty());
155 assert_eq!(tree.pop(), None);
156 }
157
158 #[test]
159 fn test_partial_fill() {
160 let mut tree = LoserTree::with_capacity(4);
161
162 tree.push(MergeEntry { item: 2, index: 0 });
163 tree.push(MergeEntry { item: 1, index: 1 });
164
165 assert_eq!(tree.len(), 2);
166
167 assert_eq!(tree.pop().unwrap().item, 1);
168 assert_eq!(tree.pop().unwrap().item, 2);
169 assert!(tree.pop().is_none());
170 }
171
172 #[test]
173 fn test_rebuild_after_pop() {
174 let mut tree = LoserTree::with_capacity(4);
175
176 tree.push(MergeEntry { item: 3, index: 0 });
177 tree.push(MergeEntry { item: 1, index: 1 });
178 tree.push(MergeEntry { item: 4, index: 2 });
179
180 assert_eq!(tree.pop().unwrap().item, 1);
181 tree.push(MergeEntry { item: 2, index: 1 });
182
183 assert_eq!(tree.pop().unwrap().item, 2);
184 assert_eq!(tree.pop().unwrap().item, 3);
185 assert_eq!(tree.pop().unwrap().item, 4);
186 }
187
188 #[test]
189 fn test_merge_tree_interface() {
190 let mut tree = LoserTree::with_capacity(4);
191
192 assert!(tree.is_empty());
194 assert_eq!(tree.len(), 0);
195 assert!(tree.pop().is_none());
196
197 tree.push(MergeEntry { item: 3, index: 0 });
199 assert_eq!(tree.len(), 1);
200 assert!(!tree.is_empty());
201
202 tree.push(MergeEntry { item: 1, index: 1 });
203 assert_eq!(tree.len(), 2);
204
205 assert_eq!(tree.pop().unwrap().item, 1);
207 assert_eq!(tree.len(), 1);
208
209 assert_eq!(tree.pop().unwrap().item, 3);
210 assert_eq!(tree.len(), 0);
211 assert!(tree.is_empty());
212 }
213
214 #[test]
215 #[should_panic(expected = "index out of bounds")]
216 fn test_push_invalid_index() {
217 let mut tree = LoserTree::with_capacity(2);
218 tree.push(MergeEntry { item: 1, index: 2 }); }
220
221 #[test]
222 fn test_clear() {
223 let mut tree = LoserTree::with_capacity(4);
224 tree.push(MergeEntry { item: 1, index: 0 });
225 tree.clear();
226 assert!(tree.is_empty());
227 assert_eq!(tree.len(), 0);
228 }
229
230 #[test]
231 fn test_capacity() {
232 let tree = LoserTree::<i32>::with_capacity(3);
233 assert_eq!(tree.capacity(), 4); }
235}