sum-segment-tree 0.1.0

A fixed-capacity sum-segment tree for weighted sampling and prioritized experience replay.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

//! Prioritized experience replay usage: store transitions in a side buffer and
//! let the sum tree sample indices proportional to their priority.

use sum_segment_tree::SumTree;

fn main() {
    let mut tree = SumTree::new(4);
    let mut data: Vec<Option<&str>> = 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);
    }

    println!("total = {}", tree.total());
    for s in [0.5, 5.5, 6.5, 8.0] {
        let (index, priority) = tree.get(s).unwrap();
        println!(
            "s={s:>4} -> index {index} = {:?} (priority {priority})",
            data[index]
        );
    }

    // Priorities change as the agent learns; update in place.
    tree.update(1, 10.0);
    println!("after update, total = {}", tree.total());
    let (index, _) = tree.get(6.0).unwrap();
    println!("s=6.0 now samples index {index} = {:?}", data[index]);
}