kxcoze

Fenwick Tree with range update and query sum

Oct 29th, 2022
1,387
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 1.06 KB | None | 0 0
  1. class Ftree:
  2.     def __init__(self, n):
  3.         self.n = n
  4.         self.tree1 = (n+1)*[0]
  5.         self.tree2 = (n+1)*[0]
  6.     def query(self, tree, r):
  7.         s = 0
  8.         while r > 0:
  9.             s += tree[r]
  10.             r -= r & -r
  11.         return s
  12.     def sum(self, r):
  13.         return self.query(self.tree1, r)*r - self.query(self.tree2, r)
  14.     def sum_range(self, l, r):
  15.         return self.sum(r) - self.sum(l-1)
  16.     def add(self, tree, k, x):
  17.         while k <= self.n:
  18.             tree[k] += x
  19.             k += k & -k
  20.     def add_range(self, l, r, x):
  21.         self.add(self.tree1, l, x)
  22.         self.add(self.tree1, r+1, -x)
  23.         self.add(self.tree2, l, x*(l-1))
  24.         self.add(self.tree2, r+1, -x *r)
  25. # Initial array
  26. a = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
  27. # Fenwick Tree
  28. t = Ftree(10)
  29. for i in range(10):
  30.     # Adding data to tree
  31.     t.add_range(i+1, i+1, a[i])
  32. # Query range sum
  33. print(t.sum_range(5, 9)) # 35
  34. # Adding value to range
  35. t.add_range(5, 9, 2)
  36.  
  37. print(t.sum_range(5, 9)) # 45
  38. t.add_range(5, 6, 1)
  39. print(t.sum_range(5, 9)) # 47
  40.  
Advertisement
Add Comment
Please, Sign In to add comment