Skip to content

Commit

Permalink
Change fenwick tree to accept uint32 to avoid extra copy
Browse files Browse the repository at this point in the history
  • Loading branch information
vmihailenco authored and caio committed Nov 4, 2018
1 parent a7e25a3 commit cbc82d5
Show file tree
Hide file tree
Showing 6 changed files with 147 additions and 20 deletions.
17 changes: 9 additions & 8 deletions Gopkg.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 0 additions & 4 deletions Gopkg.toml
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,3 @@
# [[override]]
# name = "github.com/x/y"
# version = "2.4.0"

[[constraint]]
version = "1.2.0"
name = "github.com/yourbasic/fenwick"
25 changes: 25 additions & 0 deletions internal/fenwick/LICENSE
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
BSD 2-Clause License

Copyright (c) 2017, Stefan Nilsson
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
1 change: 1 addition & 0 deletions internal/fenwick/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
This is https://github.com/yourbasic/fenwick with int64 replaced with uint32 to avoid constructing temporary []int64 slice just for the lib.
108 changes: 108 additions & 0 deletions internal/fenwick/list.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,108 @@
// 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 []uint32
}

// New creates a new list with the given elements.
func New(n ...uint32) *List {
len := len(n)
t := make([]uint32, 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) uint32 {
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 uint32) {
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 uint32) {
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) uint32 {
var sum uint32
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) uint32 {
var sum uint32
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 uint32) {
i := len(l.tree)
l.tree = append(l.tree, 0)
l.tree[i] = n - l.Get(i)
}
12 changes: 4 additions & 8 deletions summary.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ import (
"math"
"sort"

"github.com/yourbasic/fenwick"
"github.com/caio/go-tdigest/internal/fenwick"
)

type summary struct {
Expand Down Expand Up @@ -51,7 +51,7 @@ func (s *summary) Add(key float64, value uint32) error {
// Reinitialize the prefixSum cache
if s.bitree.Len() >= len(s.counts) {
for i := idx; i < len(s.counts); i++ {
s.bitree.Set(i, int64(s.counts[i]))
s.bitree.Set(i, s.counts[i])
}
} else {
s.rebuildFenwickTree()
Expand All @@ -61,11 +61,7 @@ func (s *summary) Add(key float64, value uint32) error {
}

func (s *summary) rebuildFenwickTree() {
x := make([]int64, cap(s.counts))
for i := 0; i < s.Len(); i++ {
x[i] = int64(s.counts[i])
}
s.bitree = fenwick.New(x...)
s.bitree = fenwick.New(s.counts[:cap(s.counts)]...)
}

func (s summary) Floor(x float64) int {
Expand Down Expand Up @@ -129,7 +125,7 @@ func (s *summary) setAt(index int, mean float64, count uint32) {
s.counts[index] = count
s.adjustRight(index)
s.adjustLeft(index)
s.bitree.Set(index, int64(count))
s.bitree.Set(index, count)
}

func (s *summary) adjustRight(index int) {
Expand Down

0 comments on commit cbc82d5

Please sign in to comment.