Last active
January 11, 2021 06:39
-
-
Save james4388/254e6b5ed887a75af8fc1067c64ba3dd to your computer and use it in GitHub Desktop.
Segment Tree
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| def range_sum(tree, left_child, right_child): | |
| return tree[left_child] + tree[right_child] | |
| def build_segment_tree(arr, tree=None, | |
| merge_fn=range_sum, | |
| node_index=0, lo=0, hi=None, initial=None): | |
| """ | |
| Build a segment tree based on arr of leaf node | |
| tree will have the size of 4 * len(arr) | |
| node_index is tree index | |
| lo, hi is array left and right index | |
| default merge fn is sum: this become range sum tree | |
| """ | |
| if hi is None: | |
| hi = len(arr) - 1 | |
| if tree is None: | |
| tree = [initial] * (len(arr) * 4) | |
| # Base case, hit leaf node | |
| if lo == hi: | |
| tree[node_index] = arr[lo] | |
| return | |
| mid = lo + (hi - lo) // 2 | |
| # Left child from node_index | |
| build_segment_tree(arr, tree, merge_fn, 2 * node_index + 1, lo=lo, hi=mid) | |
| # Right child from node_index | |
| build_segment_tree(arr, tree, merge_fn, 2 * node_index + 2, lo=mid+1, hi=hi) | |
| # Merge child to this current node | |
| tree[node_index] = merge_fn(tree, node_index*2+1, node_index*2+2) | |
| return tree |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment