# sum-segment-tree
A fixed-capacity sum tree for weighted sampling, the structure used in
prioritized experience replay.
Priorities are stored in the leaves of a complete binary tree packed into a flat
array. Every internal node holds the sum of its children, so the total weight
and "find the leaf at cumulative weight `s`" are both `O(log capacity)`. Writes
are a ring buffer: once the tree is full, the oldest leaf is overwritten.
## API
- `new(capacity)`: holds exactly `capacity` leaves (padded to a power of two internally).
- `push(priority)`: write a priority at the next slot, return its leaf index.
- `update(index, priority)`: change a priority in place and fix the sums.
- `total()`: sum of all priorities.
- `get(s)`: find the leaf whose cumulative interval contains `s`.
- `priority(index)`, `len()`, `capacity()`, `is_full()`.
## Usage
```rust
use sum_segment_tree::SumTree;
let mut tree = SumTree::new(4);
// Store transitions in your own buffer, keyed by the returned index.
let mut data = vec![None; tree.capacity()];
for (label, priority) in [("a", 5.0), ("b", 1.0), ("c", 1.0), ("d", 3.0)] {
let index = tree.push(priority);
data[index] = Some(label);
}
// Sample proportional to priority. With s in [0, total) the probability of a
// leaf is priority / total.
let s = 0.5 * tree.total();
let (index, priority) = tree.get(s).unwrap();
let _ = (index, priority, data[index]);
// Priorities change as the agent learns.
tree.update(1, 10.0);
```
To draw a sample, pick `s` uniformly in `[0, total())` and call `get(s)`. Leaves
with priority zero are never returned. Boundaries between intervals fall to the
earlier leaf.
## Development
```sh
cargo test # rstest table-driven tests
cargo run --example prioritized
cargo bench # criterion benchmarks for push, update, get
```