@@ -184,7 +184,68 @@ impl<K: TotalOrd, V> TreeMap<K, V> {
184
184
/// Get a lazy iterator over the key-value pairs in the map.
185
185
/// Requires that it be frozen (immutable).
186
186
pub fn iter < ' a > ( & ' a self ) -> TreeMapIterator < ' a , K , V > {
187
- TreeMapIterator { stack : ~[ ] , node : & self . root , remaining : self . length }
187
+ TreeMapIterator {
188
+ stack : ~[ ] ,
189
+ node : & self . root ,
190
+ remaining_min : self . length ,
191
+ remaining_max : self . length
192
+ }
193
+ }
194
+
195
+ /// Get a lazy iterator that should be initialized using
196
+ /// `iter_traverse_left`/`iter_traverse_right`/`iter_traverse_complete`.
197
+ fn iter_for_traversal < ' a > ( & ' a self ) -> TreeMapIterator < ' a , K , V > {
198
+ TreeMapIterator {
199
+ stack : ~[ ] ,
200
+ node : & self . root ,
201
+ remaining_min : 0 ,
202
+ remaining_max : self . length
203
+ }
204
+ }
205
+
206
+ /// Return a lazy iterator to the first key-value pair whose key is not less than `k`
207
+ /// If all keys in map are less than `k` an empty iterator is returned.
208
+ pub fn lower_bound_iter < ' a > ( & ' a self , k : & K ) -> TreeMapIterator < ' a , K , V > {
209
+ let mut iter: TreeMapIterator < ' a , K , V > = self . iter_for_traversal ( ) ;
210
+ loop {
211
+ match * iter. node {
212
+ Some ( ref r) => {
213
+ match k. cmp ( & r. key ) {
214
+ Less => iter_traverse_left ( & mut iter) ,
215
+ Greater => iter_traverse_right ( & mut iter) ,
216
+ Equal => {
217
+ iter_traverse_complete ( & mut iter) ;
218
+ return iter;
219
+ }
220
+ }
221
+ }
222
+ None => {
223
+ iter_traverse_complete ( & mut iter) ;
224
+ return iter;
225
+ }
226
+ }
227
+ }
228
+ }
229
+
230
+ /// Return a lazy iterator to the first key-value pair whose key is greater than `k`
231
+ /// If all keys in map are not greater than `k` an empty iterator is returned.
232
+ pub fn upper_bound_iter < ' a > ( & ' a self , k : & K ) -> TreeMapIterator < ' a , K , V > {
233
+ let mut iter: TreeMapIterator < ' a , K , V > = self . iter_for_traversal ( ) ;
234
+ loop {
235
+ match * iter. node {
236
+ Some ( ref r) => {
237
+ match k. cmp ( & r. key ) {
238
+ Less => iter_traverse_left ( & mut iter) ,
239
+ Greater => iter_traverse_right ( & mut iter) ,
240
+ Equal => iter_traverse_right ( & mut iter)
241
+ }
242
+ }
243
+ None => {
244
+ iter_traverse_complete ( & mut iter) ;
245
+ return iter;
246
+ }
247
+ }
248
+ }
188
249
}
189
250
190
251
/// Get a lazy iterator that consumes the treemap.
@@ -205,7 +266,8 @@ impl<K: TotalOrd, V> TreeMap<K, V> {
205
266
pub struct TreeMapIterator < ' self , K , V > {
206
267
priv stack : ~[ & ' self ~TreeNode < K , V > ] ,
207
268
priv node : & ' self Option < ~TreeNode < K , V > > ,
208
- priv remaining : uint
269
+ priv remaining_min : uint ,
270
+ priv remaining_max : uint
209
271
}
210
272
211
273
impl < ' self , K , V > Iterator < ( & ' self K , & ' self V ) > for TreeMapIterator < ' self , K , V > {
@@ -222,7 +284,10 @@ impl<'self, K, V> Iterator<(&'self K, &'self V)> for TreeMapIterator<'self, K, V
222
284
None => {
223
285
let res = self . stack . pop ( ) ;
224
286
self . node = & res. right ;
225
- self . remaining -= 1 ;
287
+ self . remaining_max -= 1 ;
288
+ if self . remaining_min > 0 {
289
+ self . remaining_min -= 1 ;
290
+ }
226
291
return Some ( ( & res. key , & res. value ) ) ;
227
292
}
228
293
}
@@ -232,7 +297,46 @@ impl<'self, K, V> Iterator<(&'self K, &'self V)> for TreeMapIterator<'self, K, V
232
297
233
298
#[ inline]
234
299
fn size_hint ( & self ) -> ( uint , Option < uint > ) {
235
- ( self . remaining , Some ( self . remaining ) )
300
+ ( self . remaining_min , Some ( self . remaining_max ) )
301
+ }
302
+ }
303
+
304
+ /// iter_traverse_left, iter_traverse_right and iter_traverse_complete are used to
305
+ /// initialize TreeMapIterator pointing to element inside tree structure.
306
+ ///
307
+ /// They should be used in following manner:
308
+ /// - create iterator using TreeMap::iter_for_traversal
309
+ /// - find required node using `iter_traverse_left`/`iter_traverse_right`
310
+ /// (current node is `TreeMapIterator::node` field)
311
+ /// - complete initialization with `iter_traverse_complete`
312
+ #[ inline]
313
+ fn iter_traverse_left < ' a , K , V > ( it : & mut TreeMapIterator < ' a , K , V > ) {
314
+ let node = it. node . get_ref ( ) ;
315
+ it. stack . push ( node) ;
316
+ it. node = & node. left ;
317
+ }
318
+
319
+ #[ inline]
320
+ fn iter_traverse_right < ' a , K , V > ( it : & mut TreeMapIterator < ' a , K , V > ) {
321
+ it. node = & ( it. node . get_ref ( ) . right ) ;
322
+ }
323
+
324
+ /// iter_traverse_left, iter_traverse_right and iter_traverse_complete are used to
325
+ /// initialize TreeMapIterator pointing to element inside tree structure.
326
+ ///
327
+ /// Completes traversal. Should be called before using iterator.
328
+ /// Iteration will start from `self.node`.
329
+ /// If `self.node` is None iteration will start from last node from which we
330
+ /// traversed left.
331
+ #[ inline]
332
+ fn iter_traverse_complete < ' a , K , V > ( it : & mut TreeMapIterator < ' a , K , V > ) {
333
+ static none: Option < ~TreeNode < K , V > > = None ;
334
+ match * it. node {
335
+ Some ( ref n) => {
336
+ it. stack . push ( n) ;
337
+ it. node = & none;
338
+ }
339
+ None => ( )
236
340
}
237
341
}
238
342
@@ -417,6 +521,20 @@ impl<T: TotalOrd> TreeSet<T> {
417
521
TreeSetIterator { iter : self . map . iter ( ) }
418
522
}
419
523
524
+ /// Get a lazy iterator pointing to the first value not less than `v` (greater or equal).
525
+ /// If all elements in the set are less than `v` empty iterator is returned.
526
+ #[ inline]
527
+ pub fn lower_bound_iter < ' a > ( & ' a self , v : & T ) -> TreeSetIterator < ' a , T > {
528
+ TreeSetIterator { iter : self . map . lower_bound_iter ( v) }
529
+ }
530
+
531
+ /// Get a lazy iterator pointing to the first value greater than `v`.
532
+ /// If all elements in the set are not greater than `v` empty iterator is returned.
533
+ #[ inline]
534
+ pub fn upper_bound_iter < ' a > ( & ' a self , v : & T ) -> TreeSetIterator < ' a , T > {
535
+ TreeSetIterator { iter : self . map . upper_bound_iter ( v) }
536
+ }
537
+
420
538
/// Visit all values in reverse order
421
539
#[ inline]
422
540
pub fn each_reverse ( & self , f : & fn ( & T ) -> bool ) -> bool {
@@ -983,6 +1101,31 @@ mod test_treemap {
983
1101
assert_eq ! ( * v, n * 2 ) ;
984
1102
n += 1 ;
985
1103
}
1104
+ assert_eq ! ( n, 5 ) ;
1105
+ }
1106
+
1107
+ #[ test]
1108
+ fn test_interval_iteration ( ) {
1109
+ let mut m = TreeMap :: new ( ) ;
1110
+ for i in range ( 1 , 100 ) {
1111
+ assert ! ( m. insert( i * 2 , i * 4 ) ) ;
1112
+ }
1113
+
1114
+ for i in range ( 1 , 198 ) {
1115
+ let mut lb_it = m. lower_bound_iter ( & i) ;
1116
+ let ( & k, & v) = lb_it. next ( ) . unwrap ( ) ;
1117
+ let lb = i + i % 2 ;
1118
+ assert_eq ! ( lb, k) ;
1119
+ assert_eq ! ( lb * 2 , v) ;
1120
+
1121
+ let mut ub_it = m. upper_bound_iter ( & i) ;
1122
+ let ( & k, & v) = ub_it. next ( ) . unwrap ( ) ;
1123
+ let ub = i + 2 - i % 2 ;
1124
+ assert_eq ! ( ub, k) ;
1125
+ assert_eq ! ( ub * 2 , v) ;
1126
+ }
1127
+ let mut end_it = m. lower_bound_iter ( & 199 ) ;
1128
+ assert_eq ! ( end_it. next( ) , None ) ;
986
1129
}
987
1130
988
1131
#[ test]
@@ -1256,7 +1399,6 @@ mod test_set {
1256
1399
1257
1400
let mut n = 0 ;
1258
1401
for x in m. iter ( ) {
1259
- printfln ! ( x) ;
1260
1402
assert_eq ! ( * x, n) ;
1261
1403
n += 1
1262
1404
}
0 commit comments