Class: FenwickTree

Inherits:
Object
  • Object
show all
Defined in:
lib/fenwick_tree.rb

Overview

Fenwick Tree

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(arg) ⇒ FenwickTree

Returns a new instance of FenwickTree.



5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
 
# File 'lib/fenwick_tree.rb', line 5

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

Instance Attribute Details

#data ⇒ Object (readonly)

Returns the value of attribute data.



3
4
5
 
# File 'lib/fenwick_tree.rb', line 3

def data
  @data
end

#size ⇒ Object (readonly)

Returns the value of attribute size.



3
4
5
 
# File 'lib/fenwick_tree.rb', line 3

def size
  @size
end

Instance Method Details

#_sum(i) ⇒ Object Also known as: left_sum 



53
54
55
56
57
58
59
60
 
# File 'lib/fenwick_tree.rb', line 53

def _sum(i)
  res = 0
  while i > 0
    res += @data[i]
    i &= i - 1
  end
  res
end

#add(i, x) ⇒ Object 



24
25
26
27
28
29
30
 
# File 'lib/fenwick_tree.rb', line 24

def add(i, x)
  i += 1
  while i <= @size
    @data[i] += x
    i += (i & -i)
  end
end

#sum(a, b = nil) ⇒ Object

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



35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
 
# File 'lib/fenwick_tree.rb', line 35

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

#to_s ⇒ Object 



63
64
65
 
# File 'lib/fenwick_tree.rb', line 63

def to_s
  "<#{self.class}: @size=#{@size}, (#{@data[1..].join(', ')})>"
end