It is a recurring and really useful idea to count some property in O(n) in an array if the relation is some kind of equation where right hand side and left hand side only differs in their indexes. For example counting number of subarrays with sum X.

1. Subarrays with sum X

  • Let’s say we are given an array A of length N and we have to count how many subarrays of A sum up to a given number X. First of all to support subarray sum query in O(1) time we need to calculate prefix sum for A. After that we have two options:

vector<int> prefix(n+1);
for(int i=1;i<=n;i++)
prefix[i]=prefix[i-1]+a[i-1];

  • Option 1. Naively set start and end of a subarray and check if its sum is X, can be done in O(n²).

for(int start=1;start<=n;start++)
  for(int end=start;end<=n;end++)
  {
    int sum=prefix[r]-prefix[l-1];
    if(sum==X)ans++;
  }
  
  return ans;
  
  • Option 2. Now as we have seen in option 1 we simply need to count all (l,r) pairs such that prefix[r] - prefix[l-1] = X or prefix[r] - X= prefix[l-1]. Now, if we store count of all prefix sums seen previously, at every index i we can query for how many prefix[i] - X we have seen before and add them to the answer as they all will be the starting indexes of subarrays ending at i. It is that simple.

unordered_map<int,int> f;
int ans=0;
for(int i=1;i<=n;i++)
{
  ans+=f[prefix[i]-X];
  f[prefix[i]]++;
}

2. Solving ABC 146 E. Rem of Sum is num

This post is work in progress. You can solve these problems till I complete it.

Some relevant problems