add RedBlackTree#select and aliases for #search and #select

CHANGELOG.md

@@ -1,5 +1,8 @@
 ## [Unreleased]
 
+- Alias `RedBlackTree#search` as `RedBlackTree#find`
+- Add `RedBlackTree#select`, aliased as `RedBlackTree#filter` and `RedBlackTree#find_all`
+
 ## [0.1.4] - 2024-10-27
 
 - Fix incorrect `super` fallback in `RedBlackTree::DataDelegation#respond_to_missing?`

lib/red-black-tree.rb

@@ -224,7 +224,7 @@ def delete! node
   # @return [RedBlackTree::Node, nil] the matching node
   def search data = nil, &block
     if block_given?
-      raise ArgumentError, "provide either data or block, not both" if data
+      raise ArgumentError, "provide either data or block, not both for search" if data
 
       _search_by_block block, @root
     else
@@ -233,6 +233,15 @@ def search data = nil, &block
       _search_by_data data, @root
     end
   end
+  alias_method :find, :search
+
+  def select &block
+    raise ArgumentError, "block must be provided for select" unless block
+
+    _select_by_block block, @root
+  end
+  alias_method :filter, :select
+  alias_method :find_all, :select
 
   # Returns true if there is a node which matches the given data/value, and false if there is not.
   #
@@ -327,14 +336,6 @@ def rotate_sub_tree! node, direction
     opp_direction_child
   end
 
-  def _search_by_block block, node
-    traverse node do |current|
-      next if current.leaf?
-
-      return current if block.call current
-    end
-  end
-
   def _search_by_data data, node
     return if node.nil? || node.leaf?
     return node if data == node.data
@@ -347,6 +348,24 @@ def _search_by_data data, node
     end
   end
 
+  def _search_by_block block, node
+    traverse node do |current|
+      next if current.leaf?
+
+      return current if block.call current
+    end
+  end
+
+  def _select_by_block block, node
+    [].tap do |result|
+      traverse node do |current|
+        next if current.leaf?
+
+        result << current if block.call current
+      end
+    end
+  end
+
   def is_root? node
     node && @root && node.object_id == @root.object_id
   end

test/test_red_black_tree.rb

@@ -699,7 +699,7 @@ def test_search_with_data_and_block
           node.data == 10
         end
       end
-      assert_equal "provide either data or block, not both", error.message
+      assert_equal "provide either data or block, not both for search", error.message
     end
 
     def test_new_tree_data_search
@@ -775,6 +775,56 @@ def test_sub_tree_block_search
     end
   end
 
+  class TestSelect < Minitest::Test
+    def test_select_without_block
+      tree = RedBlackTree.new
+      error = assert_raises ArgumentError do
+        tree.select
+      end
+      assert_equal "block must be provided for select", error.message
+    end
+
+    def test_new_tree_block_select
+      tree = RedBlackTree.new
+      result = tree.select do |node|
+        node.data == 10
+      end
+      assert_equal [], result
+    end
+
+    def test_single_node_tree_block_select
+      tree = RedBlackTree.new
+      node_10 = IntegerNode.new 10
+      tree << node_10
+      result = tree.select do |node|
+        node.data == 10
+      end
+      assert_equal [node_10], result
+    end
+
+    def test_sub_tree_block_select
+      tree = RedBlackTree.new
+      node_10 = IntegerNode.new 10
+      tree << node_10
+      node_5 = IntegerNode.new 5
+      tree << node_5
+      node_15 = IntegerNode.new 15
+      tree << node_15
+      node_1 = IntegerNode.new 1
+      tree << node_1
+      node_4 = IntegerNode.new 4
+      tree << node_4
+      result = tree.select do |node|
+        node.data % 5 == 0
+      end
+      assert_equal [node_5, node_10, node_15], result
+      result = tree.select do |node|
+        node.data == 25
+      end
+      assert_equal [], result
+    end
+  end
+
   class TestInclude < Minitest::Test
     def test_new_tree_include
       tree = RedBlackTree.new