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

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

cargo test                    # rstest table-driven tests
cargo run --example prioritized
cargo bench                   # criterion benchmarks for push, update, get