interval_tree.rb 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200
  1. #!/usr/bin/env ruby
  2. # frozen_string_literal: true
  3. # Copyright (c) 2011-2021 MISHIMA, Hiroyuki; Simeon Simeonov; Carlos Alonso; Sam Davies; amarzot-yesware; Daniel Petri Rocha
  4. # Permission is hereby granted, free of charge, to any person obtaining
  5. # a copy of this software and associated documentation files (the
  6. # "Software"), to deal in the Software without restriction, including
  7. # without limitation the rights to use, copy, modify, merge, publish,
  8. # distribute, sublicense, and/or sell copies of the Software, and to
  9. # permit persons to whom the Software is furnished to do so, subject to
  10. # the following conditions:
  11. # The above copyright notice and this permission notice shall be
  12. # included in all copies or substantial portions of the Software.
  13. # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  14. # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
  15. # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  16. # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
  17. # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
  18. # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
  19. # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
  20. module IntervalTree
  21. class Tree
  22. def initialize(ranges, &range_factory)
  23. range_factory = ->(l, r) { (l...r + 1) } unless block_given?
  24. ranges_excl = ensure_exclusive_end([ranges].flatten, range_factory)
  25. ranges_excl.sort_by! { |x| [x.begin, x.end] }
  26. @top_node = divide_intervals(ranges_excl)
  27. end
  28. attr_reader :top_node
  29. def divide_intervals(intervals)
  30. return nil if intervals.empty?
  31. x_center = center(intervals)
  32. s_center = []
  33. s_left = []
  34. s_right = []
  35. intervals.each do |k|
  36. if k.end < x_center
  37. s_left << k
  38. elsif k.begin > x_center
  39. s_right << k
  40. else
  41. s_center << k
  42. end
  43. end
  44. s_center.sort_by! { |x| [x.begin, x.end] }
  45. Node.new(x_center, s_center,
  46. divide_intervals(s_left), divide_intervals(s_right))
  47. end
  48. # Search by range or point
  49. DEFAULT_OPTIONS = { unique: true, sort: true }.freeze
  50. def search(query, options = {})
  51. options = DEFAULT_OPTIONS.merge(options)
  52. return nil unless @top_node
  53. if query.respond_to?(:begin)
  54. result = top_node.search(query)
  55. options[:unique] ? result.uniq! : result
  56. else
  57. result = point_search(top_node, query, [], options[:unique])
  58. end
  59. options[:sort] ? result.sort_by { |x| [x.begin, x.end] } : result
  60. end
  61. def ==(other)
  62. top_node == other.top_node
  63. end
  64. private
  65. def ensure_exclusive_end(ranges, range_factory)
  66. ranges.map do |range|
  67. if !range.respond_to?(:exclude_end?)
  68. range
  69. elsif range.exclude_end?
  70. range
  71. else
  72. range_factory.call(range.begin, range.end)
  73. end
  74. end
  75. end
  76. def center(intervals)
  77. (
  78. intervals.map(&:begin).min +
  79. intervals.map(&:end).max
  80. ) / 2
  81. end
  82. def point_search(node, point, result, unique)
  83. stack = [node]
  84. until stack.empty?
  85. node = stack.pop
  86. node_left_node = node.left_node
  87. node_right_node = node.right_node
  88. node_x_center = node.x_center
  89. node_s_center_end = node.s_center_end
  90. traverse_left = (point < node_x_center)
  91. if node_s_center_end && point < node_s_center_end
  92. node.s_center.each do |k|
  93. break if k.begin > point
  94. result << k if point < k.end
  95. end
  96. end
  97. if node_left_node && traverse_left
  98. stack << node_left_node
  99. elsif node_right_node && !traverse_left
  100. stack << node_right_node
  101. end
  102. end
  103. if unique
  104. result.uniq
  105. else
  106. result
  107. end
  108. end
  109. end
  110. class Node
  111. def initialize(x_center, s_center, left_node, right_node)
  112. @x_center = x_center
  113. @s_center = s_center
  114. @s_center_end = s_center.map(&:end).max
  115. @left_node = left_node
  116. @right_node = right_node
  117. end
  118. attr_reader :x_center, :s_center, :s_center_end, :left_node, :right_node
  119. def ==(other)
  120. x_center == other.x_center &&
  121. s_center == other.s_center &&
  122. left_node == other.left_node &&
  123. right_node == other.right_node
  124. end
  125. # Search by range only
  126. def search(query)
  127. search_s_center(query) +
  128. (left_node && query.begin < x_center && left_node.search(query) || []) +
  129. (right_node && query.end > x_center && right_node.search(query) || [])
  130. end
  131. private
  132. def search_s_center(query)
  133. result = []
  134. s_center.each do |k|
  135. k_begin = k.begin
  136. query_end = query.end
  137. break if k_begin > query_end
  138. k_end = k.end
  139. query_begin = query.begin
  140. k_begin_gte_q_begin = k_begin >= query_begin
  141. k_end_lte_q_end = k_end <= query_end
  142. next unless
  143. (
  144. # k is entirely contained within the query
  145. k_begin_gte_q_begin &&
  146. k_end_lte_q_end
  147. ) || (
  148. # k's start overlaps with the query
  149. k_begin_gte_q_begin &&
  150. (k_begin < query_end)
  151. ) || (
  152. # k's end overlaps with the query
  153. (k_end > query_begin) &&
  154. k_end_lte_q_end
  155. ) || (
  156. # k is bigger than the query
  157. (k_begin < query_begin) &&
  158. (k_end > query_end)
  159. )
  160. result << k
  161. end
  162. result
  163. end
  164. end
  165. end