Prefix Sum

1.1 Fenwick Tree

Keywords: Contribution Trick

1.1.1 Classic Fenwick

Just like prefix sums, we store 0 for the first index (1-based indexing).
Fenwick Trees are also called Binary Indexed Tree/BIT.

The core idea is to store interval sums in an unbalanced-tree way (known as lowbit).

It supports two operations - query an interval sum and update a single value.

Oftenly, a Fenwick Tree is faster than a Segment Tree.

int query(int index)
{
  int val=0;
  for(int i=index+1;i>0;i-=lowbit(i))
    val+=fenwick[i];
  return val;
}

void update(int index,int value)
{
  int delta=value-nums[index];
  nums[index]=value;
  for(int i=index+1;i<=n;i+=lowbit(i))
    fenwick[i]+=delta;
}

The way to create Fenwick Tree is \(\mathcal O(n)\).

for(int i=1;i<=n;++i)
{
  int par=i+lowbit(i);
  if(par<=n)
    fenwick[par]+=fenwick[i];
}

1.1.2 Interval Fenwick

We can modify a little bit to accomplish interval updates. In spite of storing the original array, store a difference array d1.

What if a range query? Use a math trick, if we want to calculate the sum of a prefix sum

\[S(x)=\sum_{i=1}^x \sum_{j=1}^i D(j)=\sum_{i=1}^x (x-i+1)D(i)=(x+1)\sum_{i=1}^x D(i)-\sum_{i=1}^x iD(i). \]

Therefore, store the \(iD(i)\) array as d2.

void updateRange(int left,int right,int value)
{
  update(d1,left,value);
  update(d1,right+1,-value);
  update(d2,left,value*left);
  update(d2,right+1,-value*(right+1));
}

int queryRange(int index)
{
  return (index+1)*prefixSum(d1,index)-prefixSum(d2,index);
}

1.2 Segment Tree

Keywords: Divide and Conquer, Lazy Propagation

We still use 1-based indexing.

The core idea is to store interval sums in a balanced-tree way (known as lazy mark).

It supports two operations - query an interval sum and update an interval value.

The efficiency comes from merging lazy marks.

void pushLazyMark(int p,int cl,int cr)
{
  int mid=(cl+cr)/2;
  mark[2*p]+=mark[p];
  segment[2*p]+=mark[p]*(mid-cl+1);
  mark[2*p+1]+=mark[p];
  segment[2*p+1]+=mark[p]*(cr-mid);
  mark[p]=0;
}

int query(int l,int r,int p,int cl,int cr)
{
  if(cl>r||cr<l)
    return 0;
  if(cl>=l&&cr<=r)
    return segment[p];
  else
  {
    pushLazyMark(p,cl,cr);
    int mid=(cl+cr)/2;
    return query(l,r,2*p,cl,mid)+query(l,r,2*p+1,mid+1,cr);
  }
}

void update(int l,int r,int val,int p,int cl,int cr)
{
  if(cr<l||cl>r)
    return;
  if(cl>=l&&cr<=r)
  {
    segment[p]+=val*(cr-cl+1);
    mark[p]+=val;
  }
  else
  {
    int mid=(cl+cr)/2;
    pushLazyMark(p,cl,cr);
    update(l,r,val,2*p,cl,mid);
    update(l,r,val,2*p+1,mid+1,cr);
    segment[p]=segment[2*p]+segment[2*p+1];
  }
}

The way to create Segment Tree is \(\mathcal O(2n)\), for there are \(2n\) nodes.

void build(int l,int r,int p)
{
  if(l==r)
  {
    segment[p]=nums[l-1];
    return;
  }
  int mid=(l+r)/2;
  build(l,mid,2*p);
  build(mid+1,r,2*p+1);
  segment[p]=segment[2*p]+segment[2*p+1];
}

Exercise

Have a rest! Now let's see a famous problem list.

1. Partition Array Into Three Parts With Equal Sum

Given an array of integers arr, return the number of ways that we can partition the array into three non-empty parts with equal sums.

Think about the question in this way: we only need to find two prefixSums, that is, t/3 and 2t/3. Store how many 2t/3 appeared after every index in a vector, so for each t/3, we can directly read how many 2t/3's after it.

2. Subarray Sum Equals K

Given an array of integers nums and an integer k, return the total number of subarrays whose sum equals to k. A subarray is a contiguous non-empty sequence of elements within an array.

Use a hashmap to collect every element seen before, so for every index, check the value minus k appears in the hashmap or not.

3. Count of Smaller Numbers After Self*

Given an integer array nums, return an integer array counts where counts[i] is the number of smaller elements to the right of nums[i].

Suppose we use a hashmap to collect every element. We will need to calculate the prefixSum of each index, so instead use a Fenwick Tree. Note we use discretization here.

class Solution
{
    int N;
    vector<int> fenwick;

    int lowbit(int x)
    {
        return x&(-x);
    }

    void add(int index,int delta)
    {
        for(int i=index+1;i<=N;i+=lowbit(i))
            fenwick[i]+=delta;
    }

    int query(int index)
    {
        int val=0;
        for(int i=index+1;i>0;i-=lowbit(i))
            val+=fenwick[i];
        return val;
    }

public:
    vector<int> countSmaller(vector<int>& nums)
    {
        int n=(int)nums.size();
        vector<int>ans(n,0);
        vector<int>discrete_nums=nums;
        sort(discrete_nums.begin(),discrete_nums.end());
        auto it=unique(discrete_nums.begin(),discrete_nums.end());
        discrete_nums.erase(it,discrete_nums.end());
        N=(int)discrete_nums.size();
        fenwick.resize(N+1,0);
        for(int i=n-1;i>=0;--i)
        {
            int rank=lower_bound(discrete_nums.begin(),discrete_nums.end(),nums[i])-discrete_nums.begin()+1;
            ans[i]=query(rank-1);
            add(rank,1);
        }
        return ans;
    }
};