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]]++;
}

1398C. Good Subarrays

  • Abridged Problem Statement : Count number of subarrays where sum of the subarray is equal to size of the subarray.

  • More formally count all pairs (l,r) such that: $$ Sum(l,r)=r-l+1 $$

  • Let $p[i]$ be sum of the prefix of length $i$ of the array. Then we can write: $$ pre[r]-pre[l-1] = r-l+1$$

  • Some manipulation of above equation leads to: $$ pre[r]-r = pre[l-1]-(l-1) $$

  • Above equation has the required property that we need to solve the problem using this technique that is right hand side and left hand side only differ in their index.

  • So we run a loop from $1$ to $n$ and for current index calculate its contribution to the answer as $r$. In simpler words just add to the answer all the occurences of $pre[r]-r$ occured so far, because all those occurences will be the leftmost indexes of the subarrays that are good.

map<int,int> m;
long long ans=0;
m[0]=1;// or start the loop from 0
for(int i=1;i<=n;i++){
	ans+=m[pre[i]-i];
	m[pre[i]-i]++;
}
cout<<ans<"\n";

ABC 146 E. Rem of sum is num

  • Abridged problem statement: Count subarrays such that remainder of its sum when divided by $K$ is equal to the size of the subarray.

  • Formally: $$Sum(l,r)\equiv (r-l+1) mod K$$

  • Or $$pre[r]-pre[l-1]\equiv r-(l-1) mod K$$

  • So if $pre[i]$ is the remainder we get when we divide sum of prefix of length $i$ by $K$.

  • Then we can write $pre[r]-pre[l-1] = (r-l+1)$ as

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

Practice Problems