segtree#

Source code: tianshou/data/utils/segtree.py

class SegmentTree(size: int)[source]#

Implementation of Segment Tree.

The segment tree stores an array arr with size n. It supports value update and fast query of the sum for the interval [left, right) in O(log n) time. The detailed procedure is as follows:

  1. Pad the array to have length of power of 2, so that leaf nodes in the segment tree have the same depth.

  2. Store the segment tree in a binary heap.

Parameters:

size – the size of segment tree.

get_prefix_sum_idx(value: float | ndarray) → int | ndarray[source]#

Find the index with given value.

Return the minimum index for each v in value so that \(v \le \mathrm{sums}_i\), where \(\mathrm{sums}_i = \sum_{j = 0}^{i} \mathrm{arr}_j\).

Warning

Please make sure all of the values inside the segment tree are non-negative when using this function.

reduce(start: int = 0, end: int | None = None) → float[source]#

Return operation(value[start:end]).