| | import operator |
| |
|
| |
|
| | class SegmentTree(object): |
| | def __init__(self, capacity, operation, neutral_element): |
| | """Build a Segment Tree data structure. |
| | |
| | https://en.wikipedia.org/wiki/Segment_tree |
| | |
| | Can be used as regular array, but with two |
| | important differences: |
| | |
| | a) setting item's value is slightly slower. |
| | It is O(lg capacity) instead of O(1). |
| | b) user has access to an efficient ( O(log segment size) ) |
| | `reduce` operation which reduces `operation` over |
| | a contiguous subsequence of items in the array. |
| | |
| | Paramters |
| | --------- |
| | capacity: int |
| | Total size of the array - must be a power of two. |
| | operation: lambda obj, obj -> obj |
| | and operation for combining elements (eg. sum, max) |
| | must form a mathematical group together with the set of |
| | possible values for array elements (i.e. be associative) |
| | neutral_element: obj |
| | neutral element for the operation above. eg. float('-inf') |
| | for max and 0 for sum. |
| | """ |
| | assert capacity > 0 and capacity & (capacity - 1) == 0, "capacity must be positive and a power of 2." |
| | self._capacity = capacity |
| | self._value = [neutral_element for _ in range(2 * capacity)] |
| | self._operation = operation |
| |
|
| | def _reduce_helper(self, start, end, node, node_start, node_end): |
| | if start == node_start and end == node_end: |
| | return self._value[node] |
| | mid = (node_start + node_end) // 2 |
| | if end <= mid: |
| | return self._reduce_helper(start, end, 2 * node, node_start, mid) |
| | else: |
| | if mid + 1 <= start: |
| | return self._reduce_helper(start, end, 2 * node + 1, mid + 1, node_end) |
| | else: |
| | return self._operation( |
| | self._reduce_helper(start, mid, 2 * node, node_start, mid), |
| | self._reduce_helper(mid + 1, end, 2 * node + 1, mid + 1, node_end) |
| | ) |
| |
|
| | def reduce(self, start=0, end=None): |
| | """Returns result of applying `self.operation` |
| | to a contiguous subsequence of the array. |
| | |
| | self.operation(arr[start], operation(arr[start+1], operation(... arr[end]))) |
| | |
| | Parameters |
| | ---------- |
| | start: int |
| | beginning of the subsequence |
| | end: int |
| | end of the subsequences |
| | |
| | Returns |
| | ------- |
| | reduced: obj |
| | result of reducing self.operation over the specified range of array elements. |
| | """ |
| | if end is None: |
| | end = self._capacity |
| | if end < 0: |
| | end += self._capacity |
| | end -= 1 |
| | return self._reduce_helper(start, end, 1, 0, self._capacity - 1) |
| |
|
| | def __setitem__(self, idx, val): |
| | |
| | idx += self._capacity |
| | self._value[idx] = val |
| | idx //= 2 |
| | while idx >= 1: |
| | self._value[idx] = self._operation( |
| | self._value[2 * idx], |
| | self._value[2 * idx + 1] |
| | ) |
| | idx //= 2 |
| |
|
| | def __getitem__(self, idx): |
| | assert 0 <= idx < self._capacity |
| | return self._value[self._capacity + idx] |
| |
|
| |
|
| | class SumSegmentTree(SegmentTree): |
| | def __init__(self, capacity): |
| | super(SumSegmentTree, self).__init__( |
| | capacity=capacity, |
| | operation=operator.add, |
| | neutral_element=0.0 |
| | ) |
| |
|
| | def sum(self, start=0, end=None): |
| | """Returns arr[start] + ... + arr[end]""" |
| | return super(SumSegmentTree, self).reduce(start, end) |
| |
|
| | def find_prefixsum_idx(self, prefixsum): |
| | """Find the highest index `i` in the array such that |
| | sum(arr[0] + arr[1] + ... + arr[i - i]) <= prefixsum |
| | |
| | if array values are probabilities, this function |
| | allows to sample indexes according to the discrete |
| | probability efficiently. |
| | |
| | Parameters |
| | ---------- |
| | perfixsum: float |
| | upperbound on the sum of array prefix |
| | |
| | Returns |
| | ------- |
| | idx: int |
| | highest index satisfying the prefixsum constraint |
| | """ |
| | assert 0 <= prefixsum <= self.sum() + 1e-5 |
| | idx = 1 |
| | while idx < self._capacity: |
| | if self._value[2 * idx] > prefixsum: |
| | idx = 2 * idx |
| | else: |
| | prefixsum -= self._value[2 * idx] |
| | idx = 2 * idx + 1 |
| | return idx - self._capacity |
| |
|
| |
|
| | class MinSegmentTree(SegmentTree): |
| | def __init__(self, capacity): |
| | super(MinSegmentTree, self).__init__( |
| | capacity=capacity, |
| | operation=min, |
| | neutral_element=float('inf') |
| | ) |
| |
|
| | def min(self, start=0, end=None): |
| | """Returns min(arr[start], ..., arr[end])""" |
| |
|
| | return super(MinSegmentTree, self).reduce(start, end) |
| |
|
