add ceilingiterator flooriterator

maps/treemap/iterator.go

@@ -14,26 +14,51 @@ var _ containers.ReverseIteratorWithKey[string, int] = (*Iterator[string, int])(
 
 // Iterator holding the iterator's state
 type Iterator[K comparable, V any] struct {
-	iterator *rbt.Iterator[K, V]
+	iterator  *rbt.Iterator[K, V]
+	rgitfirst bool
 }
 
 // Iterator returns a stateful iterator whose elements are key/value pairs.
 func (m *Map[K, V]) Iterator() *Iterator[K, V] {
 	return &Iterator[K, V]{iterator: m.tree.Iterator()}
 }
 
+func (m *Map[K, V]) FloorIterator(key K) *Iterator[K, V] {
+	it := &Iterator[K, V]{iterator: m.tree.FloorIterator(key)}
+	if it.iterator.Node() != nil {
+		it.rgitfirst = true
+	}
+	return it
+}
+
+func (m *Map[K, V]) CeilingIterator(key K) *Iterator[K, V] {
+	it := &Iterator[K, V]{iterator: m.tree.CeilingIterator(key)}
+	if it.iterator.Node() != nil {
+		it.rgitfirst = true
+	}
+	return it
+}
+
 // Next moves the iterator to the next element and returns true if there was a next element in the container.
 // If Next() returns true, then next element's key and value can be retrieved by Key() and Value().
 // If Next() was called for the first time, then it will point the iterator to the first element if it exists.
 // Modifies the state of the iterator.
 func (iterator *Iterator[K, V]) Next() bool {
+	if iterator.rgitfirst {
+		iterator.rgitfirst = false
+		return true
+	}
 	return iterator.iterator.Next()
 }
 
 // Prev moves the iterator to the previous element and returns true if there was a previous element in the container.
 // If Prev() returns true, then previous element's key and value can be retrieved by Key() and Value().
 // Modifies the state of the iterator.
 func (iterator *Iterator[K, V]) Prev() bool {
+	if iterator.rgitfirst {
+		iterator.rgitfirst = false
+		return true
+	}
 	return iterator.iterator.Prev()
 }
 

trees/avltree/iterator.go

@@ -27,6 +27,51 @@ func (tree *Tree[K, V]) Iterator() *Iterator[K, V] {
 	return &Iterator[K, V]{tree: tree, node: nil, position: begin}
 }
 
+func (tree *Tree[K, V]) FloorIerator(key K) *Iterator[K, V] {
+	found := false
+	var floor *Node[K, V]
+	n := tree.Root
+	for n != nil {
+		c := tree.Comparator(key, n.Key)
+		switch {
+		case c == 0:
+			return &Iterator[K, V]{tree: tree, node: n, position: between}
+			return n, true
+		case c < 0:
+			n = n.Children[0]
+		case c > 0:
+			floor, found = n, true
+			n = n.Children[1]
+		}
+	}
+	if found {
+		return &Iterator[K, V]{tree: tree, node: n, position: between}
+	}
+	return &Iterator[K, V]{tree: tree, node: nil, position: begin}
+}
+
+func (tree *Tree[K, V]) CeilingIterator(key K) *Iterator[K, V] {
+	found := false
+	var floor *Node[K, V]
+	n := tree.Root
+	for n != nil {
+		c := tree.Comparator(key, n.Key)
+		switch {
+		case c == 0:
+			return &Iterator[K, V]{tree: tree, node: n, position: between}
+		case c < 0:
+			floor, found = n, true
+			n = n.Children[0]
+		case c > 0:
+			n = n.Children[1]
+		}
+	}
+	if found {
+		return &Iterator[K, V]{tree: tree, node: n, position: between}
+	}
+	return &Iterator[K, V]{tree: tree, node: nil, position: end}
+}
+
 // Next moves the iterator to the next element and returns true if there was a next element in the container.
 // If Next() returns true, then next element's key and value can be retrieved by Key() and Value().
 // If Next() was called for the first time, then it will point the iterator to the first element if it exists.

trees/redblacktree/iterator.go

@@ -27,6 +27,50 @@ func (tree *Tree[K, V]) Iterator() *Iterator[K, V] {
 	return &Iterator[K, V]{tree: tree, node: nil, position: begin}
 }
 
+func (tree *Tree[K, V]) FloorIterator(key K) *Iterator[K, V] {
+	found := false
+	var floor *Node[K, V]
+	node := tree.Root
+	for node != nil {
+		compare := tree.Comparator(key, node.Key)
+		switch {
+		case compare == 0:
+			return &Iterator[K, V]{tree: tree, node: node, position: between}
+		case compare < 0:
+			node = node.Left
+		case compare > 0:
+			floor, found = node, true
+			node = node.Right
+		}
+	}
+	if found {
+		return &Iterator[K, V]{tree: tree, node: floor, position: between}
+	}
+	return &Iterator[K, V]{tree: tree, node: nil, position: begin}
+}
+
+func (tree *Tree[K, V]) CeilingIterator(key K) *Iterator[K, V] {
+	found := false
+	var ceiling *Node[K, V]
+	node := tree.Root
+	for node != nil {
+		compare := tree.Comparator(key, node.Key)
+		switch {
+		case compare == 0:
+			return &Iterator[K, V]{tree: tree, node: node, position: between}
+		case compare < 0:
+			ceiling, found = node, true
+			node = node.Left
+		case compare > 0:
+			node = node.Right
+		}
+	}
+	if found {
+		return &Iterator[K, V]{tree: tree, node: ceiling, position: between}
+	}
+	return &Iterator[K, V]{tree: tree, node: nil, position: end}
+}
+
 // IteratorAt returns a stateful iterator whose elements are key/value pairs that is initialised at a particular node.
 func (tree *Tree[K, V]) IteratorAt(node *Node[K, V]) *Iterator[K, V] {
 	return &Iterator[K, V]{tree: tree, node: node, position: between}