Range Sum Query Mutable — BIT / Segment Tree

Sanjeev SharmaSanjeev Sharma
1 min read

Advertisement

Problem

Implement update(i, val) and sumRange(l, r) with mutable array.


Approach 1 — BIT (Fenwick Tree)

Time: O(log n) update and query | Space: O(n)


Solutions

Python — BIT

class NumArray:
    def __init__(self, nums: list[int]):
        self.n=len(nums); self.bit=[0]*(self.n+1)
        for i,v in enumerate(nums): self._update(i+1,v)
    def _update(self,i,delta):
        while i<=self.n: self.bit[i]+=delta; i+=i&(-i)
    def _query(self,i):
        s=0
        while i>0: s+=self.bit[i]; i-=i&(-i)
        return s
    def update(self, index: int, val: int) -> None:
        self._update(index+1,val-self.sumRange(index,index))
    def sumRange(self, left: int, right: int) -> int:
        return self._query(right+1)-self._query(left)

Python — Segment Tree

class NumArray:
    def __init__(self, nums: list[int]):
        self.n=len(nums); self.tree=[0]*(4*self.n)
        self._build(nums,1,0,self.n-1)
    def _build(self,a,node,l,r):
        if l==r: self.tree[node]=a[l]; return
        m=(l+r)//2
        self._build(a,2*node,l,m); self._build(a,2*node+1,m+1,r)
        self.tree[node]=self.tree[2*node]+self.tree[2*node+1]
    def _update(self,node,l,r,i,v):
        if l==r: self.tree[node]=v; return
        m=(l+r)//2
        if i<=m: self._update(2*node,l,m,i,v)
        else: self._update(2*node+1,m+1,r,i,v)
        self.tree[node]=self.tree[2*node]+self.tree[2*node+1]
    def _query(self,node,l,r,ql,qr):
        if qr<l or r<ql: return 0
        if ql<=l and r<=qr: return self.tree[node]
        m=(l+r)//2
        return self._query(2*node,l,m,ql,qr)+self._query(2*node+1,m+1,r,ql,qr)
    def update(self,i,v): self._update(1,0,self.n-1,i,v)
    def sumRange(self,l,r): return self._query(1,0,self.n-1,l,r)

C++ — BIT

class NumArray {
    vector<int> bit; int n;
    void upd(int i,int d){for(;i<=n;i+=i&-i)bit[i]+=d;}
    int qry(int i){int s=0;for(;i>0;i-=i&-i)s+=bit[i];return s;}
public:
    NumArray(vector<int>& nums):n(nums.size()),bit(n+1,0){
        for(int i=0;i<n;i++)upd(i+1,nums[i]);
    }
    void update(int i,int v){upd(i+1,v-sumRange(i,i));}
    int sumRange(int l,int r){return qry(r+1)-qry(l);}
};

Complexity

  • Time: O(log n) per operation | Space: O(n)

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro