utils.segment_tree

SegmentTree

Please Reference ding/ding/utils/segment_tree.py for usage.

class ding.utils.segment_tree.SegmentTree(capacity: int, operation: Callable, neutral_element: Optional[float] = None)

Overview:

Segment tree data structure, implemented by the tree-like array. Only the leaf nodes are real value, non-leaf nodes are to do some operations on its left and right child.

Interface:

__init__, reduce, __setitem__, __getitem__

__getitem__(idx: int) → float

Overview:

Get leaf[idx]

Arguments:

• idx (int): Leaf node index(relative index), add capacity to change to absolute index.

Returns:

• val (float): The value of leaf[idx]

__init__(capacity: int, operation: Callable, neutral_element: Optional[float] = None) → None

Overview:

Initialize the segment tree. Tree’s root node is at index 1.

Arguments:

• capacity (int): Capacity of the tree (the number of the leaf nodes), should be the power of 2.

• operation (function): The operation function to construct the tree, e.g. sum, max, min, etc.

• neutral_element (float or None): The value of the neutral element, which is used to init all nodes value in the tree.

__setitem__(idx: int, val: float) → None

Overview:

Set leaf[idx] = val; Then update the related nodes.

Arguments:

• idx (int): Leaf node index(relative index), should add capacity to change to absolute index.

• val (float): The value that will be assigned to leaf[idx].

reduce(start: int = 0, end: Optional[int] = None) → float

Overview:

Reduce the tree in range [start, end)

Arguments:

• start (int): Start index(relative index, the first leaf node is 0), default set to 0

• end (int or None): End index(relative index), default set to self.capacity

Returns:

• reduce_result (float): The reduce result value, which is dependent on data type and operation

class ding.utils.segment_tree.SumSegmentTree(capacity: int)

__getitem__(idx: int) → float

Overview:

Get leaf[idx]

Arguments:

• idx (int): Leaf node index(relative index), add capacity to change to absolute index.

Returns:

• val (float): The value of leaf[idx]

__init__(capacity: int) → None

Overview:

Init sum segment tree by passing operation='sum'

__setitem__(idx: int, val: float) → None

Overview:

Set leaf[idx] = val; Then update the related nodes.

Arguments:

• idx (int): Leaf node index(relative index), should add capacity to change to absolute index.

• val (float): The value that will be assigned to leaf[idx].

find_prefixsum_idx(prefixsum: float, trust_caller: bool = True) → int

Overview:

Find the highest non-zero index i, sum_{j}leaf[j] <= prefixsum (where 0 <= j < i) and sum_{j}leaf[j] > prefixsum (where 0 <= j < i+1)

Arguments:

• prefixsum (float): The target prefixsum.

• trust_caller (bool): Whether to trust caller, which means whether to check whether this tree’s sum is greater than the input prefixsum by calling reduce function.

  Default set to True.

Returns:

• idx (int): Eligible index.

reduce(start: int = 0, end: Optional[int] = None) → float

Overview:

Reduce the tree in range [start, end)

Arguments:

• start (int): Start index(relative index, the first leaf node is 0), default set to 0

• end (int or None): End index(relative index), default set to self.capacity

Returns:

• reduce_result (float): The reduce result value, which is dependent on data type and operation

class ding.utils.segment_tree.MinSegmentTree(capacity: int)

__getitem__(idx: int) → float

Overview:

Get leaf[idx]

Arguments:

• idx (int): Leaf node index(relative index), add capacity to change to absolute index.

Returns:

• val (float): The value of leaf[idx]

__init__(capacity: int) → None

Overview:

Init sum segment tree by passing operation='min'

__setitem__(idx: int, val: float) → None

Overview:

Set leaf[idx] = val; Then update the related nodes.

Arguments:

• idx (int): Leaf node index(relative index), should add capacity to change to absolute index.

• val (float): The value that will be assigned to leaf[idx].

reduce(start: int = 0, end: Optional[int] = None) → float

Overview:

Reduce the tree in range [start, end)

Arguments:

• start (int): Start index(relative index, the first leaf node is 0), default set to 0

• end (int or None): End index(relative index), default set to self.capacity

Returns:

• reduce_result (float): The reduce result value, which is dependent on data type and operation