Advertisement
Iam_Sandeep

segment trees implementation

Sep 1st, 2021
1,145
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 2.40 KB | None | 0 0
  1. # Online Python compiler (interpreter) to run Python online.
  2. # Write Python 3 code in this online editor and run it.
  3. from math import  log2, ceil
  4. #st is segment tree
  5. #if no of nodes is perfect squre then size of seg tree array=2N-1
  6. #if no of nodes is not perfect sqaure then we should take nearest greater perfect square (lets say x). then segment tree size=2x-1
  7. def make(arr):
  8.     n=len(arr)
  9.     h=ceil(log2(n))
  10.     size=2*2**h-1
  11.     st=[None]*size
  12.     print(arr)
  13.     #l--->left pointer
  14.     #r---->right pointer
  15.     #i----->index to store the sum in seg tree array
  16.     #child 1--> stores at 2*i+1
  17.     #child 2 stores at 2*i+2
  18.     def build(i,l,r):
  19.         if l>r:
  20.             return
  21.         elif l==r:#like merge sort base condtions
  22.             st[i]=arr[l]
  23.         else:
  24.             mid=(l+r)//2
  25.             build(2*i+1,l,mid)
  26.             build(2*i+2,mid+1,r)
  27.             st[i]=st[2*i+1]+st[2*i+2]
  28.     #l,r---->constants
  29.     # we have to compute the sum of elements from index l to r including them
  30.     #so these are fixed
  31.     #s---->left pointer
  32.     #r---->right pointer
  33.     def findsum(i,l,r,s,e):
  34.         #complete over lap
  35.         if l<=s and e<=r:
  36.             return st[i]
  37.         #no overlap
  38.         elif e<l or s>r:
  39.             return 0
  40.             #partial overlap
  41.         else:
  42.             mid=(s+e)//2
  43.             one=findsum(2*i+1,l,r,s,mid)
  44.             two=findsum(2*i+2,l,r,mid+1,e)
  45.             return one+two
  46.     build(0,0,n-1)
  47.     print(st)#print segment tree
  48.     print(findsum(0,1,3,0,n-1))
  49.     #now we shall perform updates on segment tree --->O(log n)
  50.     ci=1#index where we have to change value
  51.     cv=2#value stores at the index that we have to change
  52.     diff=arr[ci]-cv#difference that we have to updated in segment tree
  53.     #make change in array first
  54.     arr[ci]=cv
  55.     def update(i,start,end):
  56.         #if the node at index i holds the cv variable in its range.Lets substract
  57.         if start<=ci<=end:
  58.             #some leaves might be None if the size(arr) is not perfect square
  59.             #lets donnt disturb it
  60.             if st[i]==None:
  61.                 return
  62.             st[i]-=diff
  63.             mid=(start+end)//2
  64.             if start==end:
  65.                 return
  66.             update(2*i+1,start,mid)
  67.             update(2*i+2,mid+1,end)
  68.     #update functon
  69.     update(0,0,n-1)
  70.     print(findsum(0,1,3,0,n-1))
  71.     print(st)
  72. make([1,3,5,7,9,11])
  73.  
  74.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement