-
Notifications
You must be signed in to change notification settings - Fork 5
Expand file tree
/
Copy pathsegment_tree.py
More file actions
executable file
·224 lines (200 loc) · 8.99 KB
/
Copy pathsegment_tree.py
File metadata and controls
executable file
·224 lines (200 loc) · 8.99 KB
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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
from functools import partial, lru_cache
from typing import Callable, Optional
import numpy as np
import ding
from .default_helper import one_time_warning
@lru_cache()
def njit():
try:
if ding.enable_numba:
import numba
from numba import njit as _njit
version = numba.__version__
middle_version = version.split(".")[1]
if int(middle_version) < 53:
_njit = partial # noqa
one_time_warning(
"Due to your numba version <= 0.53.0, DI-engine disables it. And you can install \
numba==0.53.0 if you want to speed up something"
)
else:
_njit = partial
except ImportError:
one_time_warning("If you want to use numba to speed up segment tree, please install numba first")
_njit = partial
return _njit
class SegmentTree:
"""
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__``
"""
def __init__(self, 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 (:obj:`int`): Capacity of the tree (the number of the leaf nodes), should be the power of 2.
- operation (:obj:`function`): The operation function to construct the tree, e.g. sum, max, min, etc.
- neutral_element (:obj:`float` or :obj:`None`): The value of the neutral element, which is used to init \
all nodes value in the tree.
"""
assert capacity > 0 and capacity & (capacity - 1) == 0
self.capacity = capacity
self.operation = operation
# Set neutral value(initial value) for all elements.
if neutral_element is None:
if operation == 'sum':
neutral_element = 0.
elif operation == 'min':
neutral_element = np.inf
elif operation == 'max':
neutral_element = -np.inf
else:
raise ValueError("operation argument should be in min, max, sum (built in python functions).")
self.neutral_element = neutral_element
# Index 1 is the root; Index ranging in [capacity, 2 * capacity - 1] are the leaf nodes.
# For each parent node with index i, left child is value[2*i] and right child is value[2*i+1].
self.value = np.full([capacity * 2], neutral_element)
self._compile()
def reduce(self, start: int = 0, end: Optional[int] = None) -> float:
"""
Overview:
Reduce the tree in range ``[start, end)``
Arguments:
- start (:obj:`int`): Start index(relative index, the first leaf node is 0), default set to 0
- end (:obj:`int` or :obj:`None`): End index(relative index), default set to ``self.capacity``
Returns:
- reduce_result (:obj:`float`): The reduce result value, which is dependent on data type and operation
"""
# TODO(nyz) check if directly reduce from the array(value) can be faster
if end is None:
end = self.capacity
assert (start < end)
# Change to absolute leaf index by adding capacity.
start += self.capacity
end += self.capacity
return _reduce(self.value, start, end, self.neutral_element, self.operation)
def __setitem__(self, idx: int, val: float) -> None:
"""
Overview:
Set ``leaf[idx] = val``; Then update the related nodes.
Arguments:
- idx (:obj:`int`): Leaf node index(relative index), should add ``capacity`` to change to absolute index.
- val (:obj:`float`): The value that will be assigned to ``leaf[idx]``.
"""
assert (0 <= idx < self.capacity), idx
# ``idx`` should add ``capacity`` to change to absolute index.
_setitem(self.value, idx + self.capacity, val, self.operation)
def __getitem__(self, idx: int) -> float:
"""
Overview:
Get ``leaf[idx]``
Arguments:
- idx (:obj:`int`): Leaf node ``index(relative index)``, add ``capacity`` to change to absolute index.
Returns:
- val (:obj:`float`): The value of ``leaf[idx]``
"""
assert (0 <= idx < self.capacity)
return self.value[idx + self.capacity]
def _compile(self) -> None:
f64 = np.array([0, 1], dtype=np.float64)
f32 = np.array([0, 1], dtype=np.float32)
i64 = np.array([0, 1], dtype=np.int64)
for d in [f64, f32, i64]:
_setitem(d, 0, 3.0, 'sum')
_reduce(d, 0, 1, 0.0, 'min')
_find_prefixsum_idx(d, 1, 0.5, 0.0)
class SumSegmentTree(SegmentTree):
def __init__(self, capacity: int) -> None:
"""
Overview:
Init sum segment tree by passing ``operation='sum'``
"""
super(SumSegmentTree, self).__init__(capacity, operation='sum')
def find_prefixsum_idx(self, 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 (:obj:`float`): The target prefixsum.
- trust_caller (:obj:`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 (:obj:`int`): Eligible index.
"""
if not trust_caller:
assert 0 <= prefixsum <= self.reduce() + 1e-5, prefixsum
return _find_prefixsum_idx(self.value, self.capacity, prefixsum, self.neutral_element)
class MinSegmentTree(SegmentTree):
def __init__(self, capacity: int) -> None:
"""
Overview:
Init sum segment tree by passing ``operation='min'``
"""
super(MinSegmentTree, self).__init__(capacity, operation='min')
@njit()
def _setitem(tree: np.ndarray, idx: int, val: float, operation: str) -> None:
tree[idx] = val
# Update from specified node to the root node
while idx > 1:
idx = idx >> 1 # To parent node idx
left, right = tree[2 * idx], tree[2 * idx + 1]
if operation == 'sum':
tree[idx] = left + right
elif operation == 'min':
tree[idx] = min([left, right])
@njit()
def _reduce(tree: np.ndarray, start: int, end: int, neutral_element: float, operation: str) -> float:
# Nodes in 【start, end) will be aggregated
result = neutral_element
while start < end:
if start & 1:
# If current start node (tree[start]) is a right child node, operate on start node and increase start by 1
if operation == 'sum':
result = result + tree[start]
elif operation == 'min':
result = min([result, tree[start]])
start += 1
if end & 1:
# If current end node (tree[end - 1]) is right child node, decrease end by 1 and operate on end node
end -= 1
if operation == 'sum':
result = result + tree[end]
elif operation == 'min':
result = min([result, tree[end]])
# Both start and end transform to respective parent node
start = start >> 1
end = end >> 1
return result
@njit()
def _find_prefixsum_idx(tree: np.ndarray, capacity: int, prefixsum: float, neutral_element: float) -> int:
# The function is to find a non-leaf node's index which satisfies:
# self.value[idx] > input prefixsum and self.value[idx + 1] <= input prefixsum
# In other words, we can assume that there are intervals: [num_0, num_1), [num_1, num_2), ... [num_k, num_k+1),
# the function is to find input prefixsum falls in which interval and return the interval's index.
idx = 1 # start from root node
while idx < capacity:
child_base = 2 * idx
if tree[child_base] > prefixsum:
idx = child_base
else:
prefixsum -= tree[child_base]
idx = child_base + 1
# Special case: The last element of ``self.value`` is neutral_element(0),
# and caller wants to ``find_prefixsum_idx(root_value)``.
# However, input prefixsum should be smaller than root_value.
if idx == 2 * capacity - 1 and tree[idx] == neutral_element:
tmp = idx
while tmp >= capacity and tree[tmp] == neutral_element:
tmp -= 1
if tmp != capacity:
idx = tmp
else:
raise ValueError("All elements in tree are the neutral_element(0), can't find non-zero element")
assert (tree[idx] != neutral_element)
return idx - capacity