-
Notifications
You must be signed in to change notification settings - Fork 0
/
fenwick-tree.go
76 lines (68 loc) · 1.65 KB
/
fenwick-tree.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
package datastructures
import (
"math"
"math/bits"
)
// FenwickTree represents a fenwick tree data structure.
type FenwickTree struct {
tree []int
}
// NewFenwickTree returns a new fenwick tree data structure.
func NewFenwickTree(data []int) *FenwickTree {
length := len(data) + 1 // index starts from 1
ft := &FenwickTree{
tree: make([]int, length),
}
for k, v := range data {
ft.tree[k+1] = v
}
for k, v := range ft.tree {
if k == 0 { // excluding index 0
continue
}
j := k + ft.lsb(k)
if j < length {
ft.tree[j] = ft.tree[j] + v
}
}
return ft
}
// PrefixSum returns the prefix sum for i.
func (ft *FenwickTree) PrefixSum(index int) int {
sum := 0
for index != 0 {
sum += ft.tree[index]
index = index - ft.lsb(index)
}
return sum
}
// RangeQuery performs a range query on the fenwick tree and
// return the value.
func (ft *FenwickTree) RangeQuery(i int, j int) int {
return ft.PrefixSum(j) - ft.PrefixSum(i-1)
}
// PointAdd updates items in the fenwick tree starting from the
// point p.
func (ft *FenwickTree) PointAdd(p int, val int) {
for p < ft.Size() {
ft.tree[p] = ft.tree[p] + val
p = p + ft.lsb(p)
}
}
// Size returns the size of the fenwick tree.
//
// the fenwick tree is one index larger than the array
// it was created from because indexing in the fenwick
// tree starts from 1 not 0.
func (ft *FenwickTree) Size() int {
return len(ft.tree)
}
// lsb is a helper function for retrieving the least significant
// bit of an integer.
func (ft *FenwickTree) lsb(number int) int {
trailingZeros := bits.TrailingZeros(uint(number))
if trailingZeros < 1 {
return 1
}
return int(math.Pow(2, float64(trailingZeros)))
}