fenwick.gno
6.58 Kb · 222 lines
1// Package fenwick is an order-statistic Fenwick (binary indexed) tree over
2// non-negative int64 weights, backed by avl.Tree so it can live in realm
3// state.
4//
5// It exists to answer one question in O(log N): given a random unit k drawn
6// from [0, Total()), which index owns that unit? That is the weighted-raffle
7// draw — each index's value is its ticket count, and Select maps a ticket
8// number to its owner. The naive version of that answer walks every index
9// accumulating values until the running total passes k, which is O(N) and
10// will eventually exceed a block's gas limit. Select never walks: it descends
11// the tree's binary structure, touching at most ~log2(N) nodes.
12//
13// DESIGN DECISION — negative deltas are rejected, not supported. Add panics
14// on a negative delta and Append panics on a negative initial value. This
15// structure is append-and-grow only. The reason is that Select is only
16// meaningful while prefix sums are monotonically non-decreasing: with a
17// negative value somewhere, "the index owning unit k" is not well defined and
18// the binary descent would silently return a wrong owner rather than fail.
19// Refusing the input loudly is the only honest option, so a caller that needs
20// decreasing weights gets a panic instead of a corrupt draw. The rest of the
21// structure (Add, Append, Total, ValueAt, Len) is otherwise a general,
22// reusable Fenwick tree.
23//
24// Complexity: Append O(log N), Add O(log N), Select O(log N), Total O(1),
25// ValueAt O(log N) (one avl lookup), Len O(1).
26package fenwick
27
28import (
29 "math/bits"
30 "strconv"
31
32 "gno.land/p/nt/avl/v0"
33)
34
35// keyWidth is enough decimal digits for any non-negative int64 (max is
36// 9223372036854775807, 19 digits). Keys are fixed-width and zero-padded so
37// that avl.Tree's lexicographic key order matches numeric index order; this
38// mirrors roundKey in r/splash/beacon.
39const keyWidth = 19
40
41const zeroPad = "0000000000000000000"
42
43// indexKey renders i as a fixed-width, zero-padded decimal string.
44func indexKey(i int) string {
45 s := strconv.Itoa(i)
46 if len(s) < keyWidth {
47 s = zeroPad[:keyWidth-len(s)] + s
48 }
49 return s
50}
51
52// Tree is an order-statistic Fenwick tree.
53//
54// Storage is two avl.Trees keyed by index (never a Go map or a slice, so that
55// realm state stays deterministic and persistent):
56//
57// - values holds the caller-visible value at each 0-based index.
58// - nodes holds the Fenwick partial sums, keyed by 1-based position, where
59// node j covers the value range (j-lsb(j), j].
60//
61// total is maintained incrementally so Total is O(1).
62type Tree struct {
63 values *avl.Tree
64 nodes *avl.Tree
65 n int
66 total int64
67}
68
69// NewTree returns an empty tree.
70func NewTree() *Tree {
71 return &Tree{
72 values: avl.NewTree(),
73 nodes: avl.NewTree(),
74 }
75}
76
77// Len returns the number of indices.
78func (t *Tree) Len() int {
79 return t.n
80}
81
82// Total returns the sum of all values. O(1).
83func (t *Tree) Total() int64 {
84 return t.total
85}
86
87// node reads the 1-based Fenwick node at position j. Every position in
88// [1, n] is written at Append time, so a live position is always present;
89// Has is still checked because avl.Get returns a single value and would
90// otherwise type-assert nil on a missing key.
91func (t *Tree) node(j int) int64 {
92 key := indexKey(j)
93 if !t.nodes.Has(key) {
94 return 0
95 }
96 return t.nodes.Get(key).(int64)
97}
98
99func (t *Tree) setNode(j int, v int64) {
100 t.nodes.Set(indexKey(j), v)
101}
102
103// ValueAt returns the value stored at index i. Panics if i is out of range.
104func (t *Tree) ValueAt(i int) int64 {
105 if i < 0 || i >= t.n {
106 panic("fenwick: index out of range")
107 }
108 key := indexKey(i)
109 if !t.values.Has(key) {
110 return 0
111 }
112 return t.values.Get(key).(int64)
113}
114
115// lsb returns the lowest set bit of j (j > 0).
116func lsb(j int) int {
117 return j & -j
118}
119
120// Append adds a new index at the end with the given initial value and
121// returns its 0-based index. Panics if value is negative: see the
122// append-and-grow note in the package doc.
123func (t *Tree) Append(value int64) int {
124 if value < 0 {
125 panic("fenwick: negative value; tree is append-and-grow only")
126 }
127 t.addTotal(value)
128
129 i := t.n
130 t.n++
131 t.values.Set(indexKey(i), value)
132
133 // The new 1-based position j covers (j-lsb(j), j]. Seed its node with
134 // the new value plus the nodes of the children that decompose
135 // (j-lsb(j), j-1], which is O(log N) reads and keeps the invariant
136 // without rebuilding anything.
137 j := t.n
138 sum := value
139 for c := j - 1; c > j-lsb(j); c -= lsb(c) {
140 sum += t.node(c)
141 }
142 t.setNode(j, sum)
143 return i
144}
145
146// Add increases the value at index i by delta. i must be in [0, Len()).
147// Panics on a negative delta: this structure is append-and-grow only, and
148// silently accepting a decrease would make Select return wrong owners rather
149// than fail (see the package doc). A zero delta is accepted and is a no-op.
150func (t *Tree) Add(i int, delta int64) {
151 if i < 0 || i >= t.n {
152 panic("fenwick: index out of range")
153 }
154 if delta < 0 {
155 panic("fenwick: negative delta; tree is append-and-grow only")
156 }
157 if delta == 0 {
158 return
159 }
160 t.addTotal(delta)
161
162 key := indexKey(i)
163 t.values.Set(key, t.ValueAt(i)+delta)
164
165 for j := i + 1; j <= t.n; j += lsb(j) {
166 t.setNode(j, t.node(j)+delta)
167 }
168}
169
170// addTotal accumulates into total, refusing to wrap. An int64 overflow here
171// would corrupt every later Select, so it is a panic and not a silent wrap.
172func (t *Tree) addTotal(delta int64) {
173 if t.total+delta < t.total {
174 panic("fenwick: total overflows int64")
175 }
176 t.total += delta
177}
178
179// Prefix returns the sum of values at indices [0, i). Panics if i is outside
180// [0, Len()]. O(log N).
181func (t *Tree) Prefix(i int) int64 {
182 if i < 0 || i > t.n {
183 panic("fenwick: index out of range")
184 }
185 var sum int64
186 for j := i; j > 0; j -= lsb(j) {
187 sum += t.node(j)
188 }
189 return sum
190}
191
192// Select returns the index owning cumulative unit k, that is the unique i
193// with Prefix(i) <= k < Prefix(i+1). Panics if k is outside [0, Total()).
194// O(log N): a binary descent over the tree's node positions, never a walk
195// over indices.
196//
197// An index whose value is 0 owns no units and is therefore never returned:
198// the descent takes the largest position whose prefix is still <= k, which
199// steps past any zero-valued index.
200func (t *Tree) Select(k int64) int {
201 if k < 0 || k >= t.total {
202 panic("fenwick: k out of range")
203 }
204
205 // Start from the highest power of two <= n and halve each step.
206 pos := 0
207 rem := k
208 for step := 1 << (bits.Len(uint(t.n)) - 1); step > 0; step >>= 1 {
209 next := pos + step
210 if next > t.n {
211 continue
212 }
213 v := t.node(next)
214 if v <= rem {
215 pos = next
216 rem -= v
217 }
218 }
219 // pos is the count of indices whose cumulative sum is still <= k, which
220 // is exactly the 0-based index owning unit k.
221 return pos
222}
