Skip to content

Instantly share code, notes, and snippets.

@samacs
Created July 1, 2014 07:00
Show Gist options
  • Save samacs/fb6d92034efb81a480de to your computer and use it in GitHub Desktop.
Save samacs/fb6d92034efb81a480de to your computer and use it in GitHub Desktop.
class SegmentTree
class Segment
attr_reader :range, :value
def initialize range, value
raise ArgumentError, 'Range expected, %s given' % range.class.name unless range.is_a?(Range)
@range, @value = range, value
end
def <=> other
case cmp = @range.begin <=> other.range.begin
when 0 then @range.end <=> other.range.end
else cmp
end
end
end
def initialize data, sorted = false
@segments = case data
when Hash, Array, Enumerable then
data.collect { |range, value| Segment.new(range, value) }
else raise ArgumentError, '2-dimensional Array or Hash expected'
end
@segments.sort! unless sorted
end
def find x
return nil if x.nil?
low = 0
high = @segments.size - 1
while low <= high
mid = (low + high) / 2
case matches?(x, low, high, mid)
when -1 then high = mid - 1
when 1 then low = mid + 1
when 0 then return @segments[mid]
else return nil
end
end
nil
end
def inspect
if @segments.size > 0
"SegmentTree(#{@segments.first.range.begin}..#{@segments.last.range.end})"
else
"SegmentTree(empty)"
end
end
private
def matches?(x, low_idx, high_idx, idx)
low, high = @segments[low_idx], @segments[high_idx]
segment = @segments[idx]
left = idx > 0 && @segments[idx - 1]
right = idx < @segments.size - 1 && @segments[idx + 1]
case
when left && (low.range.begin..left.range.end).include?(x) then -1
when segment.range.include?(x) then 0
when right && (right.range.begin..high.range.end).include?(x) then 1
else nil
end
end
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment