class AcLibraryRb::FenwickTree

Fenwick Tree

Attributes

data[R]
size[R]

Public Class Methods

new(arg) click to toggle source
# File lib_lock/ac-library-rb/fenwick_tree.rb, line 6
def initialize(arg)
  case arg
  when Array
    @size = arg.size
    @data = [0].concat(arg)
    (1 ... @size).each do |i|
      up = i + (i & -i)
      next if up > @size

      @data[up] += @data[i]
    end
  when Integer
    @size = arg
    @data = Array.new(@size + 1, 0)
  else
    raise ArgumentError.new("wrong argument. type is Array or Integer")
  end
end

Public Instance Methods

_sum(i) click to toggle source
# File lib_lock/ac-library-rb/fenwick_tree.rb, line 54
def _sum(i)
  res = 0
  while i > 0
    res += @data[i]
    i &= i - 1
  end
  res
end
Also aliased as: left_sum
add(i, x) click to toggle source
# File lib_lock/ac-library-rb/fenwick_tree.rb, line 25
def add(i, x)
  i += 1
  while i <= @size
    @data[i] += x
    i += (i & -i)
  end
end
left_sum(i)
Alias for: _sum
sum(a, b = nil) click to toggle source

.sum(l, r) # [l, r) <- Original .sumĀ® # [0, r) <- [Experimental] .sum(l..r) # [l, r] <- [Experimental]

# File lib_lock/ac-library-rb/fenwick_tree.rb, line 36
def sum(a, b = nil)
  if b
    _sum(b) - _sum(a)
  elsif a.is_a?(Range)
    l = a.begin
    l += @size if l < 0
    if r = a.end
      r += @size if r < 0
      r += 1 unless a.exclude_end?
    else
      r = @size
    end
    _sum(r) - _sum(l)
  else
    _sum(a)
  end
end