add RedBlackTree#search

CHANGELOG.md

@@ -1,6 +1,7 @@
 ## [Unreleased]
 
 - Alias `RedBlackTree#left_most_node` as `RedBlackTree#min`
+- Add `RedBlackTree#search`
 
 ## [0.1.1] - 2024-08-04
 

lib/red-black-tree.rb

@@ -219,12 +219,18 @@ def delete! node
     self
   end
 
-  private
+  # Searches for a node which matches the given data/value.
+  #
+  # @param data [any] the data to search for
+  # @return [RedBlackTree::Node, nil] the matching node
+  def search data
+    raise ArgumentError, "data must be provided for search" unless data
 
-  def is_root? node
-    node.object_id == @root.object_id
+    _search data, @root
   end
 
+  private
+
   # Rotates a (sub-)tree starting from the given node in the given direction.
   #
   # @param node [RedBlackTree::Node] the root node of the sub-tree
@@ -251,6 +257,22 @@ def rotate_sub_tree! node, direction
     opp_direction_child
   end
 
+  def _search data, current_node
+    return if current_node.nil? || current_node.leaf?
+    return current_node if data == current_node.data
+
+    mock_node = current_node.class.new data
+    if mock_node >= current_node
+      _search data, current_node.right
+    else
+      _search data, current_node.left
+    end
+  end
+
+  def is_root? node
+    node.object_id == @root.object_id
+  end
+
   def increment_size!
     @size += 1
   end

test/test_red_black_tree.rb

@@ -93,7 +93,7 @@ def test_single_node_tree_insert_without_direction
       assert_equal error.message, "Target parent direction must be provided"
     end
 
-    def test_tree_with_sub_tree_insert
+    def test_sub_tree_insert
       tree = RedBlackTree.new
       node_root = RedBlackTree::StringNode.new "root"
       tree.insert! node_root
@@ -129,7 +129,7 @@ def test_tree_with_sub_tree_insert
       assert_equal node_left_child, tree.left_most_node
     end
 
-    def test_tree_with_sub_tree_insert_with_existing_child
+    def test_sub_tree_insert_with_existing_child
       tree = RedBlackTree.new
       node_root = RedBlackTree::StringNode.new "root"
       tree.insert! node_root
@@ -155,13 +155,6 @@ def test_tree_with_sub_tree_insert_with_existing_child
       end
       assert_equal error.message, "Target parent already has left child"
     end
-
-    # fine for these tests where just the == operator is required
-    class RedBlackTree::StringNode < RedBlackTree::Node
-      def <=> other
-        self.data <=> other.data
-      end
-    end
   end
 
   class TestAdd < Minitest::Test
@@ -682,6 +675,40 @@ def test_delete_non_root_black_node_with_leaf_children_and_close_nephew
     end
   end
 
+  class TestSearch < Minitest::Test
+    def test_new_tree_search
+      tree = RedBlackTree.new
+      result = tree.search 10
+      assert_nil result
+    end
+
+    def test_single_node_tree_search
+      tree = RedBlackTree.new
+      node_10 = RedBlackTree::IntegerNode.new 10
+      tree << node_10
+      result = tree.search 10
+      assert_equal node_10, result
+    end
+
+    def test_sub_tree_search
+      tree = RedBlackTree.new
+      node_10 = RedBlackTree::IntegerNode.new 10
+      tree << node_10
+      node_5 = RedBlackTree::IntegerNode.new 5
+      tree << node_5
+      node_15 = RedBlackTree::IntegerNode.new 15
+      tree << node_15
+      node_1 = RedBlackTree::IntegerNode.new 1
+      tree << node_1
+      node_4 = RedBlackTree::IntegerNode.new 4
+      tree << node_4
+      result = tree.search 5
+      assert_equal node_5, result
+      result = tree.search 15
+      assert_equal node_15, result
+    end
+  end
+
   class TestShift < Minitest::Test
     def test_new_tree_shift
       tree = RedBlackTree.new
@@ -705,7 +732,7 @@ def test_single_node_tree_shift
       assert_nil tree.left_most_node
     end
 
-    def test_tree_with_sub_tree_shift
+    def test_sub_tree_shift
       tree = RedBlackTree.new
       node_10 = RedBlackTree::IntegerNode.new 10
       tree << node_10
@@ -747,8 +774,8 @@ def test_new_tree
 
     def test_single_node_tree
       tree = RedBlackTree.new
-      node = RedBlackTree::StringNode.new "root"
-      tree.insert! node
+      node_10 = RedBlackTree::IntegerNode.new 10
+      tree.insert! node_10
       assert_equal 1, tree.size
       refute tree.empty?
     end
@@ -763,8 +790,8 @@ def test_new_tree
 
     def test_single_node_tree
       tree = RedBlackTree.new
-      node = RedBlackTree::StringNode.new "root"
-      tree.insert! node
+      node_10 = RedBlackTree::IntegerNode.new 10
+      tree.insert! node_10
       assert_equal 1, tree.size
       assert tree.any?
     end
@@ -775,4 +802,10 @@ def <=> other
       self.data <=> other.data
     end
   end
+
+  class RedBlackTree::StringNode < RedBlackTree::Node
+    def <=> other
+      self.data.length <=> other.data.length
+    end
+  end
 end