Skip to content

Commit

Permalink
Use a fenwick tree to optimize prefix sums
Browse files Browse the repository at this point in the history
This patch makes summary use the excellent fenwick package from
https://github.com/yourbasic/fenwick to dramatically speed
up Add() calls while adding a bounded cost to memory.

    benchmark             old ns/op     new ns/op     delta
    BenchmarkAdd1-4       187           206           +10.16%
    BenchmarkAdd10-4      332           274           -17.47%
    BenchmarkAdd100-4     1092          325           -70.24%

The tree creation could be faster, but profiling has shown that
it doesn't seem to be that much of a problem (that's because
`summary.setAt()` happens a *lot* more than `summary.Add()`)
  • Loading branch information
caio committed Oct 29, 2017
1 parent a124661 commit d89c9f2
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 6 deletions.
8 changes: 7 additions & 1 deletion Gopkg.lock

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

4 changes: 4 additions & 0 deletions Gopkg.toml
Original file line number Diff line number Diff line change
Expand Up @@ -24,3 +24,7 @@
[[constraint]]
revision = "5344a9259b21627d94279721ab1f27eb029194e7"
name = "github.com/leesper/go_rng"

[[constraint]]
version = "1.2.0"
name = "github.com/yourbasic/fenwick"
25 changes: 20 additions & 5 deletions summary.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,18 +4,23 @@ import (
"fmt"
"math"
"sort"

"github.com/yourbasic/fenwick"
)

type summary struct {
means []float64
counts []uint32
bitree *fenwick.List
}

func newSummary(initialCapacity uint) *summary {
return &summary{
s := &summary{
means: make([]float64, 0, initialCapacity),
counts: make([]uint32, 0, initialCapacity),
}
s.rebuildFenwickTree()
return s
}

func (s summary) Len() int {
Expand Down Expand Up @@ -43,9 +48,21 @@ func (s *summary) Add(key float64, value uint32) error {
s.means[idx] = key
s.counts[idx] = value

// Reinitialize the prefixSum cache
// we can likely be smarter when doing this
s.rebuildFenwickTree()

return nil
}

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

func (s summary) Floor(x float64) int {
return sort.Search(len(s.means), func(i int) bool {
return s.means[i] >= x
Expand All @@ -60,10 +77,7 @@ func (s summary) FindInsertionIndex(x float64) int {
}

func (s summary) HeadSum(index int) (sum float64) {
for i := 0; i < index; i++ {
sum += float64(s.counts[i])
}
return sum
return float64(s.bitree.Sum(index))
}

func (s summary) FindIndex(x float64) int {
Expand Down Expand Up @@ -110,6 +124,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))
}

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

0 comments on commit d89c9f2

Please sign in to comment.