Skip to content
Permalink
Branch: master
Find file Copy path
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
109 lines (99 sloc) 2.64 KB
// Package fenwick provides a list data structure supporting prefix sums.
//
// A Fenwick tree, or binary indexed tree, is a space-efficient list
// data structure that can efficiently update elements and calculate
// prefix sums in a list of numbers.
//
// Compared to a common array, a Fenwick tree achieves better balance
// between element update and prefix sum calculation – both operations
// run in O(log n) time – while using the same amount of memory.
// This is achieved by representing the list as an implicit tree,
// where the value of each node is the sum of the numbers in that
// subtree.
//
package fenwick
// List represents a list of numbers with support for efficient
// prefix sum computation. The zero value is an empty list.
type List struct {
// The tree slice stores range sums of an underlying array t.
// To compute the prefix sum t[0] + t[1] + t[k-1], add elements
// which correspond to each 1 bit in the binary expansion of k.
//
// For example, this is how the sum of the 13 first elements
// in t is computed: 13 is 1101₂ in binary, so the elements
// at indices 1101₂ - 1, 1100₂ - 1, and 1000₂ - 1 are added;
// they contain the range sums t[12], t[8] + … t[11], and
// t[0] + … + t[7], respectively.
//
tree []int64
}
// New creates a new list with the given elements.
func New(n ...int64) *List {
len := len(n)
t := make([]int64, len)
copy(t, n)
for i := range t {
if j := i | (i + 1); j < len {
t[j] += t[i]
}
}
return &List{
tree: t,
}
}
// Len returns the number of elements in the list.
func (l *List) Len() int {
return len(l.tree)
}
// Get returns the element at index i.
func (l *List) Get(i int) int64 {
sum := l.tree[i]
j := i + 1
j -= j & -j
for i > j {
sum -= l.tree[i-1]
i -= i & -i
}
return sum
}
// Set sets the element at index i to n.
func (l *List) Set(i int, n int64) {
n -= l.Get(i)
for len := len(l.tree); i < len; i |= i + 1 {
l.tree[i] += n
}
}
// Add adds n to the element at index i.
func (l *List) Add(i int, n int64) {
for len := len(l.tree); i < len; i |= i + 1 {
l.tree[i] += n
}
}
// Sum returns the sum of the elements from index 0 to index i-1.
func (l *List) Sum(i int) int64 {
var sum int64
for i > 0 {
sum += l.tree[i-1]
i -= i & -i
}
return sum
}
// SumRange returns the sum of the elements from index i to index j-1.
func (l *List) SumRange(i, j int) int64 {
var sum int64
for j > i {
sum += l.tree[j-1]
j -= j & -j
}
for i > j {
sum -= l.tree[i-1]
i -= i & -i
}
return sum
}
// Append appends a new element to the end of the list.
func (l *List) Append(n int64) {
i := len(l.tree)
l.tree = append(l.tree, 0)
l.tree[i] = n - l.Get(i)
}
You can’t perform that action at this time.