diff --git a/CHANGELOG.md b/CHANGELOG.md index ad2ac85..05ae229 100644 --- a/CHANGELOG.md +++ b/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?` diff --git a/lib/red-black-tree.rb b/lib/red-black-tree.rb index a8e6c26..467e77a 100644 --- a/lib/red-black-tree.rb +++ b/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 diff --git a/test/test_red_black_tree.rb b/test/test_red_black_tree.rb index b863f7d..8cbbaac 100644 --- a/test/test_red_black_tree.rb +++ b/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