segment_tree.rb

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