| 2 | |
| 3 | |
| 4 | class SegmentTree: |
| 5 | def __init__(self, a): |
| 6 | self.A = a |
| 7 | self.N = len(self.A) |
| 8 | self.st = [0] * ( |
| 9 | 4 * self.N |
| 10 | ) # approximate the overall size of segment tree with array N |
| 11 | if self.N: |
| 12 | self.build(1, 0, self.N - 1) |
| 13 | |
| 14 | def left(self, idx): |
| 15 | """ |
| 16 | Returns the left child index for a given index in a binary tree. |
| 17 | |
| 18 | >>> s = SegmentTree([1, 2, 3]) |
| 19 | >>> s.left(1) |
| 20 | 2 |
| 21 | >>> s.left(2) |
| 22 | 4 |
| 23 | """ |
| 24 | return idx * 2 |
| 25 | |
| 26 | def right(self, idx): |
| 27 | """ |
| 28 | Returns the right child index for a given index in a binary tree. |
| 29 | |
| 30 | >>> s = SegmentTree([1, 2, 3]) |
| 31 | >>> s.right(1) |
| 32 | 3 |
| 33 | >>> s.right(2) |
| 34 | 5 |
| 35 | """ |
| 36 | return idx * 2 + 1 |
| 37 | |
| 38 | def build(self, idx, left, right): |
| 39 | if left == right: |
| 40 | self.st[idx] = self.A[left] |
| 41 | else: |
| 42 | mid = (left + right) // 2 |
| 43 | self.build(self.left(idx), left, mid) |
| 44 | self.build(self.right(idx), mid + 1, right) |
| 45 | self.st[idx] = max(self.st[self.left(idx)], self.st[self.right(idx)]) |
| 46 | |
| 47 | def update(self, a, b, val): |
| 48 | """ |
| 49 | Update the values in the segment tree in the range [a,b] with the given value. |
| 50 | |
| 51 | >>> s = SegmentTree([1, 2, 3, 4, 5]) |
| 52 | >>> s.update(2, 4, 10) |
| 53 | True |
| 54 | >>> s.query(1, 5) |
| 55 | 10 |
| 56 | """ |
| 57 | return self.update_recursive(1, 0, self.N - 1, a - 1, b - 1, val) |
| 58 | |
| 59 | def update_recursive(self, idx, left, right, a, b, val): |
| 60 | """ |
| 61 | update(1, 1, N, a, b, v) for update val v to [a,b] |