make RedBlackTree#left_most_node public, add RedBlackTree#traverse{_*…

…} methods, extend RedBlackTree#search to accept a block

CHANGELOG.md

@@ -1,5 +1,13 @@
 ## [Unreleased]
 
+- Make `RedBlackTree#left_most_node` public
+- Add `RedBlackTree#traverse_pre_order`
+- Add `RedBlackTree#traverse_in_order`
+- Add `RedBlackTree#traverse_post_order`
+- Add `RedBlackTree#traverse_level_order`
+- Add `RedBlackTree#traverse`, alias of `RedBlackTree#traverse_in_order`
+- Extend `RedBlackTree#search` to accept a block
+
 ## [0.1.2] - 2024-09-08
 
 - Fix a bunch of issues in `RedBlackTree#insert!` and `RedBlackTree#delete!` algorithms

README.md

@@ -185,11 +185,6 @@ end
 
 ## WIP Features
 
-- `RedBlackTree#traverse_in_order`
-- `RedBlackTree#traverse_pre_order`
-- `RedBlackTree#traverse_post_order`
-- `RedBlackTree#traverse_level_order`
-- `RedBlackTree#traverse` (default in-order)
 - `RedBlackTree#max`
 - `RedBlackTree#height`
 - `RedBlackTree#clear`

lib/red-black-tree.rb

@@ -14,7 +14,6 @@ class RedBlackTree
   # @return [RedBlackTree::Node, nil] the root node
   attr_reader :root
 
-  # @private
   # @return [RedBlackTree::Node, nil] the left most node
   attr_reader :left_most_node
   alias_method :min, :left_most_node
@@ -220,12 +219,19 @@ def delete! node
 
   # Searches for a node which matches the given data/value.
   #
-  # @param data [any] the data to search for
+  # @param data [any, nil] the data to search for
+  # @yield [RedBlackTree::Node] the block to be used for comparison
   # @return [RedBlackTree::Node, nil] the matching node
-  def search data
-    raise ArgumentError, "data must be provided for search" unless data
+  def search data = nil, &block
+    if block_given?
+      raise ArgumentError, "provide either data or block, not both" if data
+
+      _search_by_block block, @root
+    else
+      raise ArgumentError, "data must be provided for search" unless data
 
-    _search data, @root
+      _search_by_data data, @root
+    end
   end
 
   # Returns true if there is a node which matches the given data/value, and false if there is not.
@@ -235,6 +241,65 @@ def include? data
     !!search(data)
   end
 
+  # Traverses the tree in pre-order and yields each node.
+  #
+  # @param node [RedBlackTree::Node] the node to start the traversal from
+  # @yield [RedBlackTree::Node] the block to be executed for each node
+  # @return [void]
+  def traverse_pre_order(node = @root, &block)
+    return if node.nil? || node.leaf?
+
+    block.call node
+    traverse_pre_order node.left, &block
+    traverse_pre_order node.right, &block
+  end
+
+  # Traverses the tree in in-order and yields each node.
+  #
+  # @param node [RedBlackTree::Node] the node to start the traversal from
+  # @yield [RedBlackTree::Node] the block to be executed for each node
+  # @return [void]
+  def traverse_in_order node = @root, &block
+    return if node.nil? || node.leaf?
+
+    traverse_in_order node.left, &block
+    block.call node
+    traverse_in_order node.right, &block
+  end
+  alias_method :traverse, :traverse_in_order
+
+  # Traverses the tree in post-order and yields each node.
+  #
+  # @param node [RedBlackTree::Node] the node to start the traversal from
+  # @yield [RedBlackTree::Node] the block to be executed for each node
+  # @return [void]
+  def traverse_post_order(node = @root, &block)
+    return if node.nil? || node.leaf?
+
+    traverse_post_order node.left, &block
+    traverse_post_order node.right, &block
+    block.call node
+  end
+
+  # Traverses the tree in level-order and yields each node.
+  #
+  # @param node [RedBlackTree::Node] the node to start the traversal from
+  # @yield [RedBlackTree::Node] the block to be executed for each node
+  # @return [void]
+  def traverse_level_order(&block)
+    return if @root.nil?
+
+    queue = [@root]
+    until queue.empty?
+      node = queue.shift
+      next if node.nil? || node.leaf?
+
+      block.call node
+      queue << node.left unless node.left.nil?
+      queue << node.right unless node.right.nil?
+    end
+  end
+
   private
 
   # Rotates a (sub-)tree starting from the given node in the given direction.
@@ -263,15 +328,23 @@ 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
+  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
 
-    mock_node = current_node.class.new data
-    if mock_node >= current_node
-      _search data, current_node.right
+    mock_node = node.class.new data
+    if mock_node >= node
+      _search_by_data data, node.right
     else
-      _search data, current_node.left
+      _search_by_data data, node.left
     end
   end
 

lib/red_black_tree/node.rb

@@ -2,6 +2,7 @@
 
 require_relative "utils"
 require_relative "node/leaf_node_comparable"
+require_relative "node/data_delegation"
 require_relative "node/left_right_element_referencers"
 
 class RedBlackTree
@@ -13,6 +14,7 @@ def inherited subclass
     end
 
     include Comparable
+    include DataDelegation
 
     # @return [any] the data/value representing the node
     attr_reader :data

lib/red_black_tree/node/data_delegation.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+class RedBlackTree
+  module DataDelegation
+    def method_missing(method_name, *args, &block)
+      if @data.respond_to?(method_name)
+        @data.public_send(method_name, *args, &block)
+      else
+        super
+      end
+    end
+
+    def respond_to_missing?(method_name, include_private = false)
+      @data.respond_to?(method_name, include_private) || super
+    end
+  end
+end

test/test_red_black_tree.rb

@@ -676,21 +676,47 @@ def test_delete_non_root_black_node_with_leaf_children_and_close_nephew
   end
 
   class TestSearch < Minitest::Test
-    def test_new_tree_search
+    def test_search_without_data_and_without_block
+      tree = RedBlackTree.new
+      error = assert_raises ArgumentError do
+        tree.search
+      end
+      assert_equal "data must be provided for search", error.message
+    end
+
+    def test_search_with_nil_data_and_without_block
+      tree = RedBlackTree.new
+      error = assert_raises ArgumentError do
+        tree.search nil
+      end
+      assert_equal "data must be provided for search", error.message
+    end
+
+    def test_search_with_data_and_block
+      tree = RedBlackTree.new
+      error = assert_raises ArgumentError do
+        tree.search 10 do |node|
+          node.data == 10
+        end
+      end
+      assert_equal "provide either data or block, not both", error.message
+    end
+
+    def test_new_tree_data_search
       tree = RedBlackTree.new
       result = tree.search 10
       assert_nil result
     end
 
-    def test_single_node_tree_search
+    def test_single_node_tree_data_search
       tree = RedBlackTree.new
       node_10 = IntegerNode.new 10
       tree << node_10
       result = tree.search 10
       assert_equal node_10, result
     end
 
-    def test_sub_tree_search
+    def test_sub_tree_data_search
       tree = RedBlackTree.new
       node_10 = IntegerNode.new 10
       tree << node_10
@@ -707,6 +733,46 @@ def test_sub_tree_search
       result = tree.search 15
       assert_equal node_15, result
     end
+
+    def test_new_tree_block_search
+      tree = RedBlackTree.new
+      result = tree.search do |node|
+        node.data == 10
+      end
+      assert_nil result
+    end
+
+    def test_single_node_tree_block_search
+      tree = RedBlackTree.new
+      node_10 = IntegerNode.new 10
+      tree << node_10
+      result = tree.search do |node|
+        node.data == 10
+      end
+      assert_equal node_10, result
+    end
+
+    def test_sub_tree_block_search
+      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.search do |node|
+        node.data == 5
+      end
+      assert_equal node_5, result
+      result = tree.search do |node|
+        node.data == 15
+      end
+      assert_equal node_15, result
+    end
   end
 
   class TestInclude < Minitest::Test
@@ -741,6 +807,102 @@ def test_sub_tree_include
     end
   end
 
+  class TestTraverseInPreOrder < Minitest::Test
+    def test_order
+      tree = RedBlackTree.new
+      node_10 = IntegerNode.new 10
+      tree << node_10
+      node_5 = IntegerNode.new 7
+      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
+      node_4 = IntegerNode.new 5
+      tree << node_4
+      expected_order = [10, 4, 1, 7, 5, 15]
+      actual_order = []
+      tree.traverse_pre_order do |node|
+        actual_order << node.data
+      end
+      assert_equal expected_order, actual_order
+    end
+  end
+
+  class TestTraverseInInOrder < Minitest::Test
+    def test_order
+      tree = RedBlackTree.new
+      node_10 = IntegerNode.new 10
+      tree << node_10
+      node_5 = IntegerNode.new 7
+      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
+      node_4 = IntegerNode.new 5
+      tree << node_4
+      expected_order = [1, 4, 5, 7, 10, 15]
+      actual_order = []
+      tree.traverse_in_order do |node|
+        actual_order << node.data
+      end
+      assert_equal expected_order, actual_order
+    end
+  end
+
+  class TestTraverseInPostOrder < Minitest::Test
+    def test_order
+      tree = RedBlackTree.new
+      node_10 = IntegerNode.new 10
+      tree << node_10
+      node_5 = IntegerNode.new 7
+      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
+      node_4 = IntegerNode.new 5
+      tree << node_4
+      expected_order = [1, 5, 7, 4, 15, 10]
+      actual_order = []
+      tree.traverse_post_order do |node|
+        actual_order << node.data
+      end
+      assert_equal expected_order, actual_order
+    end
+  end
+
+  class TestTraverseInLevelOrder < Minitest::Test
+    def test_order
+      tree = RedBlackTree.new
+      node_10 = IntegerNode.new 10
+      tree << node_10
+      node_5 = IntegerNode.new 7
+      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
+      node_4 = IntegerNode.new 5
+      tree << node_4
+      expected_order = [10, 4, 15, 1, 7, 5]
+      actual_order = []
+      tree.traverse_level_order do |node|
+        actual_order << node.data
+      end
+      assert_equal expected_order, actual_order
+    end
+  end
+
   class TestShift < Minitest::Test
     def test_new_tree_shift
       tree = RedBlackTree.new