fenwick.gno

  1// Package fenwick is a Binary Indexed Tree (Peter Fenwick, 1994): an array of
  2// int64 that answers a prefix sum and applies a point update in O(log n), with
  3// no extra storage beyond the array itself.
  4//
  5// The tradeoff it occupies is the whole reason to reach for one. A plain slice
  6// updates in O(1) and sums in O(n). A slice of running totals sums in O(1) and
  7// updates in O(n). A Fenwick tree does both in O(log n), which is what a realm
  8// wants when reads and writes are interleaved: a scoreboard written by every
  9// player and read by every render.
 10//
 11// The operation that actually justifies the package is [Tree.SearchPrefix]: the
 12// smallest index whose prefix sum exceeds a target, in O(log n), by descending
 13// the tree rather than scanning it. That is the primitive behind a
 14// stake-weighted draw (pick a number below Total, ask which holder owns it) and
 15// behind a leaderboard's rank lookup. Prefix sums on their own would not be
 16// worth a data structure.
 17//
 18// Reads are total and writes are strict, deliberately: see the package README.
 19//
 20// A live demo of this package is at
 21// [r/moul/x/daily/fenwickdemo](/r/moul/x/daily/fenwickdemo/v0).
 22package fenwick
 23
 24// MaxSize bounds the number of slots so a caller cannot size the tree from
 25// unbounded user input and make the allocation itself the attack. It is not a
 26// gas figure: a query costs O(log n) regardless, and the cap exists so the one
 27// O(n) operation in the package, construction, stays bounded.
 28const MaxSize = 65536
 29
 30// Tree is a Binary Indexed Tree over int64 values, addressed 0-based.
 31//
 32// The zero value is not usable; build one with [New] or [FromSlice].
 33type Tree struct {
 34	n int
 35	// t is 1-based internally: t[i] holds the sum of the LowBit(i) values
 36	// ending at i. Index 0 is unused, which is what makes the bit tricks work.
 37	t []int64
 38}
 39
 40// New returns a Tree of n zero-valued slots.
 41//
 42// It panics if n is negative or above [MaxSize]: both are programmer errors
 43// that a caller cannot recover from, and a silently clamped size would hand
 44// back a tree that disagrees with the caller about its own length.
 45func New(n int) *Tree {
 46	if n < 0 {
 47		panic("fenwick: negative size")
 48	}
 49	if n > MaxSize {
 50		panic("fenwick: size above MaxSize")
 51	}
 52	return &Tree{n: n, t: make([]int64, n+1)}
 53}
 54
 55// FromSlice returns a Tree holding vals, built in O(n) rather than by n
 56// successive Add calls, which would cost O(n log n).
 57//
 58// The slice is copied; later writes to vals do not reach the tree.
 59func FromSlice(vals []int64) *Tree {
 60	tr := New(len(vals))
 61	// Seed each slot, then push it into its parent once. Every slot is visited
 62	// exactly once, which is the whole O(n) construction.
 63	for i, v := range vals {
 64		tr.t[i+1] += v
 65		if parent := (i + 1) + lowBit(i+1); parent <= tr.n {
 66			tr.t[parent] += tr.t[i+1]
 67		}
 68	}
 69	return tr
 70}
 71
 72// Len returns the number of slots.
 73func (tr *Tree) Len() int { return tr.n }
 74
 75// Add adds delta to the value at i.
 76//
 77// It panics if i is out of range. A write is a transaction: failing loudly is
 78// the correct outcome, and the caller still has the chance to abort.
 79func (tr *Tree) Add(i int, delta int64) {
 80	if i < 0 || i >= tr.n {
 81		panic("fenwick: index out of range")
 82	}
 83	for j := i + 1; j <= tr.n; j += lowBit(j) {
 84		tr.t[j] += delta
 85	}
 86}
 87
 88// Set replaces the value at i, returning the value it displaced.
 89//
 90// It panics if i is out of range, for the same reason [Tree.Add] does.
 91func (tr *Tree) Set(i int, v int64) int64 {
 92	old := tr.At(i)
 93	if i < 0 || i >= tr.n {
 94		panic("fenwick: index out of range")
 95	}
 96	tr.Add(i, v-old)
 97	return old
 98}
 99
100// Prefix returns the sum of the first i values, that is of [0, i).
101//
102// It is total: i is clamped into [0, Len], so Prefix(-5) is 0 and Prefix(1e9)
103// is Total. Reads are what Render calls, and a panicking Render is a page that
104// can never be displayed again, on a path that can never be redeployed.
105func (tr *Tree) Prefix(i int) int64 {
106	if i > tr.n {
107		i = tr.n
108	}
109	var sum int64
110	for ; i > 0; i -= lowBit(i) {
111		sum += tr.t[i]
112	}
113	return sum
114}
115
116// Range returns the sum of [lo, hi). It is total: bounds are clamped, and an
117// empty or inverted range is 0.
118func (tr *Tree) Range(lo, hi int) int64 {
119	if lo < 0 {
120		lo = 0
121	}
122	if hi > tr.n {
123		hi = tr.n
124	}
125	if lo >= hi {
126		return 0
127	}
128	return tr.Prefix(hi) - tr.Prefix(lo)
129}
130
131// At returns the value at i, or 0 if i is out of range. Total, like the other
132// reads.
133func (tr *Tree) At(i int) int64 {
134	if i < 0 || i >= tr.n {
135		return 0
136	}
137	return tr.Range(i, i+1)
138}
139
140// Total returns the sum of every value.
141func (tr *Tree) Total() int64 { return tr.Prefix(tr.n) }
142
143// Slice materializes the tree back into a plain slice, in O(n log n).
144//
145// It exists for rendering and for tests. A hot path should not call it.
146func (tr *Tree) Slice() []int64 {
147	out := make([]int64, tr.n)
148	for i := 0; i < tr.n; i++ {
149		out[i] = tr.At(i)
150	}
151	return out
152}
153
154// SearchPrefix returns the smallest index i for which Prefix(i+1) > target,
155// or Len if no prefix exceeds it. It runs in O(log n) by descending the tree.
156//
157// With non-negative values this is the weighted select: draw a target in
158// [0, Total) and SearchPrefix names the slot that owns it, each slot chosen in
159// proportion to its value. A zero-valued slot can never be selected, which is
160// the property that makes it usable for a stake-weighted draw.
161//
162// It is total, and a negative target returns 0.
163//
164// The result is only meaningful when every value is non-negative: with a
165// negative value in the tree the prefix sums stop increasing and "the smallest
166// index whose prefix exceeds target" stops being well defined. The package does
167// not forbid negative values, because Add with a negative delta is the ordinary
168// way to decrement a counter; it is SearchPrefix alone that needs them absent.
169func (tr *Tree) SearchPrefix(target int64) int {
170	if target < 0 {
171		return 0
172	}
173	pos := 0
174	// Walk the powers of two downward, largest first. At each step, take the
175	// jump if the value it would accumulate still does not exceed target.
176	for step := highBit(tr.n); step > 0; step >>= 1 {
177		if next := pos + step; next <= tr.n && tr.t[next] <= target {
178			pos = next
179			target -= tr.t[next]
180		}
181	}
182	return pos
183}
184
185// lowBit returns the lowest set bit of i, the range width the slot covers.
186func lowBit(i int) int { return i & (-i) }
187
188// highBit returns the highest power of two that is <= n, or 0 when n is 0.
189func highBit(n int) int {
190	p := 1
191	for p<<1 <= n {
192		p <<= 1
193	}
194	if n == 0 {
195		return 0
196	}
197	return p
198}
199
200// Covers returns the half-open range [lo, hi) of 0-based slots that the
201// internal node at 1-based position i summarizes, which is what makes the
202// structure legible rather than magic. It is exported for the demo and for
203// anyone trying to see the shape; it is not needed to use the tree.
204//
205// It returns (0, 0) when i is outside [1, Len].
206func (tr *Tree) Covers(i int) (int, int) {
207	if i < 1 || i > tr.n {
208		return 0, 0
209	}
210	return i - lowBit(i), i
211}