diff --git a/sorted_map/map.mbt b/sorted_map/map.mbt index 0251a10e2..3f45ef8bb 100644 --- a/sorted_map/map.mbt +++ b/sorted_map/map.mbt @@ -241,6 +241,17 @@ pub impl[K : @quickcheck.Arbitrary + Compare, V : @quickcheck.Arbitrary] @quickc @quickcheck.Arbitrary::arbitrary(size, rs) |> T::from_iter } +///|Returns a new array of key-value pairs that are within the specified range [low, high). +pub fn range[K : Compare, V]( + self : T[K, V], + low : K, + high : K +) -> Array[(K, V)] { + let result = [] + self.each(fn(k, v) { if k >= low && k <= high { result.push((k, v)) } }) + result +} + // AVL tree operations ///| diff --git a/sorted_map/map_test.mbt b/sorted_map/map_test.mbt index 29f74c9c0..cd3e4c0cf 100644 --- a/sorted_map/map_test.mbt +++ b/sorted_map/map_test.mbt @@ -211,3 +211,14 @@ test "iter" { } inspect!(buf, content="[1a][2b][3c]") } + +test "range" { + let map = @sorted_map.of([(1, "a"), (2, "b"), (3, "c"), (4, "d"), (5, "e")]) + let result = map.range(2, 4) + inspect!( + result, + content= + #|[(2, "b"), (3, "c"), (4, "d")] + , + ) +} diff --git a/sorted_map/sorted_map.mbti b/sorted_map/sorted_map.mbti index b74bafd00..8e18d6f21 100644 --- a/sorted_map/sorted_map.mbti +++ b/sorted_map/sorted_map.mbti @@ -24,6 +24,7 @@ impl T { op_equal[K : Eq, V : Eq](Self[K, V], Self[K, V]) -> Bool op_get[K : Compare, V](Self[K, V], K) -> V? op_set[K : Compare, V](Self[K, V], K, V) -> Unit + range[K : Compare, V](Self[K, V], K, K) -> Array[(K, V)] remove[K : Compare, V](Self[K, V], K) -> Unit size[K, V](Self[K, V]) -> Int to_array[K, V](Self[K, V]) -> Array[(K, V)] diff --git a/sorted_set/set.mbt b/sorted_set/set.mbt index 3e8479ec2..a7c3120a2 100644 --- a/sorted_set/set.mbt +++ b/sorted_set/set.mbt @@ -450,6 +450,13 @@ fn to_array[T](self : Node[T]) -> Array[T] { aux(Some(self)) } +///| +pub fn range[V : Compare](self : T[V], low : V, high : V) -> Array[V] { + let result = [] + self.each(fn(x) { if x >= low && x <= high { result.push(x) } }) + result +} + // AVL tree operations ///| diff --git a/sorted_set/set_test.mbt b/sorted_set/set_test.mbt index a220e4e15..7cb9344ca 100644 --- a/sorted_set/set_test.mbt +++ b/sorted_set/set_test.mbt @@ -222,3 +222,9 @@ test "from_iter empty iter" { let pq : @sorted_set.T[Int] = @sorted_set.T::from_iter(Iter::empty()) inspect!(pq, content="@sorted_set.of([])") } + +test "range" { + let set = @sorted_set.of([1, 2, 3, 4, 5]) + let result = set.range(2, 4) + inspect!(result, content="[2, 3, 4]") +} diff --git a/sorted_set/sorted_set.mbti b/sorted_set/sorted_set.mbti index 631e7cc15..a8af5cbcb 100644 --- a/sorted_set/sorted_set.mbti +++ b/sorted_set/sorted_set.mbti @@ -26,6 +26,7 @@ impl T { is_empty[V : Compare](Self[V]) -> Bool iter[V](Self[V]) -> Iter[V] op_equal[V : Compare](Self[V], Self[V]) -> Bool + range[V : Compare](Self[V], V, V) -> Array[V] remove[V : Compare](Self[V], V) -> Unit size[V : Compare](Self[V]) -> Int64 subset[V : Compare](Self[V], Self[V]) -> Bool