diff --git a/lib/interval_tree.rb b/lib/interval_tree.rb index f711bff..53dba04 100644 --- a/lib/interval_tree.rb +++ b/lib/interval_tree.rb @@ -73,17 +73,24 @@ def center(intervals) ) / 2 end - def point_search(node, point, result, unique = true) - node.s_center.each do |k| - if k.include?(point) - result << k + def point_search(node, point, result, unique) + stack = [node] + + until stack.empty? + node = stack.pop + + node.s_center.each do |k| + if k.begin <= point && point < k.end + result << k + end end - end - if node.left_node && ( point.to_r < node.x_center ) - point_search(node.left_node, point, []).each{|k|result << k} - end - if node.right_node && ( point.to_r >= node.x_center ) - point_search(node.right_node, point, []).each{|k|result << k} + if node.left_node && ( point.to_r < node.x_center ) + stack << node.left_node + + elsif node.right_node && ( point.to_r >= node.x_center ) + stack << node.right_node + end + end if unique result.uniq diff --git a/spec/interval_tree_spec.rb b/spec/interval_tree_spec.rb index f225617..669066f 100644 --- a/spec/interval_tree_spec.rb +++ b/spec/interval_tree_spec.rb @@ -276,9 +276,18 @@ def initialize(l, r, value = nil) end end - context 'with unique defaulting to true' do - context 'given intervals with duplicates' do - it 'returns the duplicates in the result' do + context 'with unique defaulting to true, given intervals with duplicates' do + context 'given [(1...3), (3...5), (3...9), (4...8), (4...8)] and a query by (3...5)]' do + it 'removes the duplicates in the result' do + itv = [(1...3), (3...5), (3...9), (4...8), (4...8)] + t = IntervalTree::Tree.new(itv) + results = t.search((3...5)) + expect(results).to match_array([(3...5), (3...9), (4...8)]) + end + end + + context 'given [(0...3), (1...4), (3...5), (0...3)] and a query by (2)' do + it 'removes the duplicates in the result' do itv = [(0...3), (1...4), (3...5), (0...3)] t = IntervalTree::Tree.new(itv) results = t.search(2) @@ -303,6 +312,31 @@ def initialize(l, r, value = nil) expect(results).to match_array([(2...4)]) end end + + context 'given [(0...5), (1...5), (3...5), (3...5)] and a query by (3)' do + it 'returns [(0...5), (1...5), (3...5), (3...5)]' do + itvs = [(0...5), (1...5), (3...5), (3...5)] + results = IntervalTree::Tree.new(itvs).search(3, unique: false) + expect(results).to match_array([(0...5), (1...5), (3...5), (3...5)]) + end + end + + context 'given [(0...3), (1...4), (3...4), (3...4), (3...5)] and a query by (3)' do + it 'returns [(1...4), (3...4), (3...4), (3...5)]' do + itvs = [(0...3), (1...4), (3...4), (3...4), (3...5)] + results = IntervalTree::Tree.new(itvs).search(3, unique: false) + expect(results).to match_array([(1...4), (3...4), (3...4), (3...5)]) + end + end + + context 'given [(0...2), (0...2), (1...2), (1...2), (2...5)] and a query by (1)' do + it 'returns [(0...2), (0...2), (1...2), (1...2)]' do + itvs = [(0...2), (0...2), (1...2), (1...2), (2...5)] + results = IntervalTree::Tree.new(itvs).search(1, unique: false) + expect(results).to match_array([(0...2), (0...2), (1...2), (1...2)]) + end + end + end context "when concerned with performance" do