Sorts s_center to improve search

README.md

@@ -44,11 +44,12 @@ See spec/interval_tree_spec.rb for details.
 ```ruby
 require "interval_tree"
 
-itv = [(0...3), (1...4), (3...5), (0...3)]
+itv = [(0...3), (1...4), (2...10), (3...5), (0...3)]
 t = IntervalTree::Tree.new(itv)
-p t.search(2)     #=> [0...3, 1...4]
-p t.search(2, unique: false) #=> [0...3, 0...3, 1...4]
-p t.search(1...4) #=> [0...3, 1...4, 3...5]
+p t.search(2)     #=> [0...3, 1...4, 2...10]
+p t.search(2, unique: false) #=> [0...3, 0...3, 1...4, 2...10]
+p t.search(2, unique: false, sort: false) #=> [2...10, 0...3, 0...3, 1...4]
+p t.search(1...4) #=> [0...3, 1...4, 2...10, 3...5]
 ```
 
 ## Note

lib/interval_tree.rb

@@ -1,50 +1,75 @@
 #!/usr/bin/env ruby
+# frozen_string_literal: true
 
-module IntervalTree
+# Copyright (c) 2011-2021 MISHIMA, Hiroyuki; Simeon Simeonov; Carlos Alonso; Sam Davies; amarzot-yesware; Daniel Petri Rocha
+
+# Permission is hereby granted, free of charge, to any person obtaining
+# a copy of this software and associated documentation files (the
+# "Software"), to deal in the Software without restriction, including
+# without limitation the rights to use, copy, modify, merge, publish,
+# distribute, sublicense, and/or sell copies of the Software, and to
+# permit persons to whom the Software is furnished to do so, subject to
+# the following conditions:
+
+# The above copyright notice and this permission notice shall be
+# included in all copies or substantial portions of the Software.
 
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
+# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+module IntervalTree
   class Tree
     def initialize(ranges, &range_factory)
-      range_factory = lambda { |l, r| (l ... r+1) } unless block_given?
+      range_factory = ->(l, r) { (l...r + 1) } unless block_given?
       ranges_excl = ensure_exclusive_end([ranges].flatten, range_factory)
+      ranges_excl.sort_by! { |x| [x.begin, x.end] }
       @top_node = divide_intervals(ranges_excl)
     end
     attr_reader :top_node
 
     def divide_intervals(intervals)
       return nil if intervals.empty?
+
       x_center = center(intervals)
-      s_center = Array.new
-      s_left = Array.new
-      s_right = Array.new
+      s_center = []
+      s_left = []
+      s_right = []
 
       intervals.each do |k|
-        case
-        when k.last.to_r < x_center
+        if k.end.to_r < x_center
           s_left << k
-        when k.first.to_r > x_center
+        elsif k.begin.to_r > x_center
           s_right << k
         else
           s_center << k
         end
       end
+
+      s_center.sort_by! { |x| [x.begin, x.end] }
+
       Node.new(x_center, s_center,
                divide_intervals(s_left), divide_intervals(s_right))
     end
 
     # Search by range or point
-    DEFAULT_OPTIONS = {unique: true}
+    DEFAULT_OPTIONS = { unique: true, sort: true }.freeze
     def search(query, options = {})
       options = DEFAULT_OPTIONS.merge(options)
 
       return nil unless @top_node
 
-      if query.respond_to?(:first)
+      if query.respond_to?(:begin)
         result = top_node.search(query)
-        options[:unique] ? result.uniq : result
+        options[:unique] ? result.uniq! : result
       else
-        point_search(self.top_node, query, [], options[:unique])
+        result = point_search(top_node, query, [], options[:unique])
       end
-        .sort_by{|x|[x.first, x.last]}
+      options[:sort] ? result.sort_by { |x| [x.begin, x.end] } : result
     end
 
     def ==(other)
@@ -55,13 +80,12 @@ def ==(other)
 
     def ensure_exclusive_end(ranges, range_factory)
       ranges.map do |range|
-        case
-        when !range.respond_to?(:exclude_end?)
+        if !range.respond_to?(:exclude_end?)
           range
-        when range.exclude_end?
+        elsif range.exclude_end?
           range
         else
-          range_factory.call(range.first, range.end)
+          range_factory.call(range.begin, range.end)
         end
       end
     end
@@ -73,72 +97,104 @@ 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_left_node = node.left_node
+        node_right_node = node.right_node
+        node_x_center = node.x_center
+        node_s_center_end = node.s_center_end
+        traverse_left = (point < node_x_center)
+
+        if node_s_center_end && point < node_s_center_end
+          node.s_center.each do |k|
+            break if k.begin > point
+
+            result << k if point < k.end
+          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 && traverse_left
+          stack << node_left_node
+
+        elsif node_right_node && !traverse_left
+          stack << node_right_node
+        end
+
       end
       if unique
         result.uniq
       else
         result
       end
     end
-  end # class Tree
+  end
 
   class Node
     def initialize(x_center, s_center, left_node, right_node)
       @x_center = x_center
       @s_center = s_center
+      @s_center_end = s_center.map(&:end).max
       @left_node = left_node
       @right_node = right_node
     end
-    attr_reader :x_center, :s_center, :left_node, :right_node
+    attr_reader :x_center, :s_center, :s_center_end, :left_node, :right_node
 
     def ==(other)
       x_center == other.x_center &&
-      s_center == other.s_center &&
-      left_node == other.left_node &&
-      right_node == other.right_node
+        s_center == other.s_center &&
+        left_node == other.left_node &&
+        right_node == other.right_node
     end
 
     # Search by range only
     def search(query)
       search_s_center(query) +
-        (left_node && query.begin.to_r < x_center && left_node.search(query) || []) +
-        (right_node && query.end.to_r > x_center && right_node.search(query) || [])
+        (left_node && query.begin < x_center && left_node.search(query) || []) +
+        (right_node && query.end > x_center && right_node.search(query) || [])
     end
 
     private
 
     def search_s_center(query)
-      s_center.select do |k|
+      result = []
+
+      s_center.each do |k|
+        k_begin = k.begin
+        query_end = query.end
+
+        break if k_begin > query_end
+
+        k_end = k.end
+        query_begin = query.begin
+
+        k_begin_gte_q_begin = k_begin >= query_begin
+        k_end_lte_q_end = k_end <= query_end
+        next unless
         (
           # k is entirely contained within the query
-          (k.begin >= query.begin) &&
-          (k.end <= query.end)
+          k_begin_gte_q_begin &&
+          k_end_lte_q_end
         ) || (
           # k's start overlaps with the query
-          (k.begin >= query.begin) &&
-          (k.begin < query.end)
+          k_begin_gte_q_begin &&
+          (k_begin < query_end)
         ) || (
           # k's end overlaps with the query
-          (k.end > query.begin) &&
-          (k.end <= query.end)
+          (k_end > query_begin) &&
+          k_end_lte_q_end
         ) || (
           # k is bigger than the query
-          (k.begin < query.begin) &&
-          (k.end > query.end)
+          (k_begin < query_begin) &&
+          (k_end > query_end)
         )
+
+        result << k
       end
-    end
-  end # class Node
 
-end # module IntervalTree
+      result
+    end
+  end
+end