From 3d50490517789402f2aebd88a4ef2d63bb37f5a7 Mon Sep 17 00:00:00 2001 From: epowsal Date: Mon, 4 Nov 2024 20:26:46 +0800 Subject: [PATCH] add ceilingiterator flooriterator --- maps/treemap/iterator.go | 27 +++++++++++++++++++- trees/avltree/iterator.go | 45 ++++++++++++++++++++++++++++++++++ trees/redblacktree/iterator.go | 44 +++++++++++++++++++++++++++++++++ 3 files changed, 115 insertions(+), 1 deletion(-) diff --git a/maps/treemap/iterator.go b/maps/treemap/iterator.go index 9bdf91b5..5fd4b881 100644 --- a/maps/treemap/iterator.go +++ b/maps/treemap/iterator.go @@ -14,7 +14,8 @@ 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. @@ -22,11 +23,31 @@ 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() } @@ -34,6 +55,10 @@ func (iterator *Iterator[K, V]) Next() bool { // 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() } diff --git a/trees/avltree/iterator.go b/trees/avltree/iterator.go index 8541ce0a..10f6f435 100644 --- a/trees/avltree/iterator.go +++ b/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. diff --git a/trees/redblacktree/iterator.go b/trees/redblacktree/iterator.go index 9aa5605e..b6526bc9 100644 --- a/trees/redblacktree/iterator.go +++ b/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}