diff --git a/pkg/treemap/custom.go b/pkg/treemap/custom.go index 7ca03e76..4e2d2691 100644 --- a/pkg/treemap/custom.go +++ b/pkg/treemap/custom.go @@ -89,3 +89,15 @@ func (m *TreeMap[K, V]) Range(f func(key K, value V) bool) { } } } + +// ReverseRange calls f sequentially in reverse order for each key and value +// present in the map. If f returns false, range stops the iteration. +func (m *TreeMap[K, V]) ReverseRange(f func(key K, value V) bool) { + iterator := m.Iterator() + iterator.End() + for iterator.Prev() { + if !f(iterator.Key(), iterator.Value()) { + break + } + } +} diff --git a/pkg/treemap/custom_test.go b/pkg/treemap/custom_test.go new file mode 100644 index 00000000..4e0c76ce --- /dev/null +++ b/pkg/treemap/custom_test.go @@ -0,0 +1,35 @@ +package treemap + +import "testing" + +func TestMapRange(t *testing.T) { + m := New[int, string]() + + m.Put(5, "e") + m.Put(6, "f") + m.Put(7, "g") + m.Put(3, "c") + m.Put(4, "d") + m.Put(1, "x") + m.Put(2, "b") + + expectedKey := 1 + m.Range(func(key int, value string) bool { + if key != expectedKey { + t.Errorf("[Range] expected %d, got %d", expectedKey, key) + return false + } + expectedKey++ + return true + }) + + expectedKey = 7 + m.ReverseRange(func(key int, value string) bool { + if key != expectedKey { + t.Errorf("[ReverseRange] expected %d, got %d", expectedKey, key) + return false + } + expectedKey-- + return true + }) +}