toolbox_rs/
loser_tree.rs

1use crate::{merge_entry::MergeEntry, merge_tree::MergeTree};
2
3pub struct LoserTree<T>
4where
5    T: Ord,
6{
7    /// array of loser indices
8    losers: Vec<usize>,
9    /// array of leaves
10    leaves: Vec<Option<MergeEntry<T>>>,
11    /// index of the winner
12    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    /// play a match between two leaves and return the index of the winner
33    #[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    /// rebuild only the path from leaf i to the root
48    fn rebuild_path(&mut self, mut i: usize) {
49        let n = self.leaves.len();
50        let internal_nodes = n - 1;
51
52        // convert leaf index to internal node index
53        i += internal_nodes;
54
55        // walk up the tree till the root
56        while i > 0 {
57            // find parent node
58            let parent = (i - 1) / 2;
59
60            // sibling of the current node either to the left or to the right
61            let sibling = if i % 2 == 0 { i - 1 } else { i + 1 };
62
63            // determine the winner of the match
64            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); // O(log n) operation
110        }
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); // next higher power of 2
131        assert_eq!(tree.losers.len(), 3); // internal nodes
132        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        // items shall be sorted
149        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        // Test empty state
193        assert!(tree.is_empty());
194        assert_eq!(tree.len(), 0);
195        assert!(tree.pop().is_none());
196
197        // Test pushing elements
198        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        // Test popping elements in order
206        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 }); // index out of bounds
219    }
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); // nächste Zweierpotenz
234    }
235}