Last active
January 29, 2024 03:53
-
-
Save qubard/a305ffd7bbc217d80f669ce28cfc3ea2 to your computer and use it in GitHub Desktop.
Fenwick Tree W/ Range And Point Updates (Python)
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
| class FenwickTreeRangeUpdatePointQuery: | |
| def __init__(self, values): | |
| if values is None: | |
| raise ValueError("Values list cannot be None!") | |
| self.N = len(values) | |
| values[0] = 0 | |
| fenwick_tree = values.copy() | |
| self.A = values | |
| # update the parents on initialization to contain | |
| # valid prefix sum values for the starting values | |
| for i in range(1, self.N): | |
| parent = i + self.lsb(i) | |
| if parent < self.N: | |
| fenwick_tree[parent] += fenwick_tree[i] | |
| self.current_tree = fenwick_tree.copy() | |
| def update_range(self, left, right, val): | |
| self.add(left, val) | |
| self.add(right + 1, -val) | |
| def add(self, i, v): | |
| while i < self.N: | |
| self.current_tree[i] += v | |
| i += self.lsb(i) | |
| # sets value at index i to v | |
| def update(self, i, v): | |
| diff = v - self.A[i] | |
| self.A[i] = v | |
| while i < self.N: | |
| self.current_tree[i] += diff | |
| i += self.lsb(i) | |
| def get(self, i): | |
| return self.prefix_sum(i, self.current_tree) - self.prefix_sum(i - 1, self.current_tree) | |
| def get_range(self, l, r): | |
| return self.prefix_sum(r, self.current_tree) - self.prefix_sum(l - 1, self.current_tree) | |
| def prefix_sum(self, i, tree): | |
| _sum = 0 | |
| while i > 0: | |
| _sum += tree[i] | |
| i -= self.lsb(i) | |
| return _sum | |
| @staticmethod | |
| def lsb(i): | |
| return i & -i | |
| # Usage | |
| n = [0] * (len(nums) + 1) | |
| for i in range(len(nums)): | |
| n[i + 1] = nums[i] | |
| self.f = FenwickTreeRangeUpdatePointQuery(n) |
Author
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
This is converted from Java re William Fiset's implementation but I have since bugfixed the line in
get()because mutations never happen against original tree, therefore
get()won't work. Not sure why that's in the original implementation unless you assumecurrentTreeis not a shallow clone oforiginalTree(which it was)