CSES Solutions – Subarray Sums II

Given an array arr[] of N integers, your task is to count the number of subarrays having sum X.

Examples:

Input: N = 5, X = 7, arr[] = {2, -1, 3, 5, -2}
Output: 2
Explanation: There are 2 subarrays with sum = 7.

Subarray {-1, 3, 5}, sum = -1 + 3 + 5 = 7.
Subarray {2, -1, 3, 5, 2}, sum = 2 - 1 + 3 + 5 + 2 = 7.
Input: N = 5, X = 3, arr[] = {1, 1, 1, -1, 1}
Output: 2
Explanation: There are 2 subarrays with sum = 3.

Subarray {1, 1, 1}, sum = 1 + 1 + 1 = 3.
Subarray {1, 1, 1, -1, 1}, sum = 1 + 1 + 1 - 1 + 1 = 3.

Approach: To solve the problem, follow the below idea:

The problem can be solved using Prefix Sums and Map. We can maintain a Map, that stores the prefix sums along with the number of times they have occurred. At each index i, let's say the prefix sum is P, that is sum of subarray arr[0...i] = P. Now if there is an index j (j < i) such that sum of arr[0...j] = P - X, then the sum of subarray arr[j+1...i] will be equal to X. Therefore, for every index i, if we can find the count of prefixes before i which have sum = P - X, we can add the count to our answer and the sum of count for all indices will be the final answer.

This allows us to find a subarray within our current array (by removing a prefix from our current prefix) that sums up to X. Also, as we iterate through the array, we continuously update the map with the new prefix sum after each step so that all possible prefix sums are counted in the map as we traverse the array.

Step-by-step algorithm:

Maintain a map, say prefSums to store the count of occurrences of each prefix sum.
Maintain a variable pref = 0, to calculate the prefix sum till any index and a variable cnt = 0 to count the number of subarrays with sum = X.
Initialize prefSums[0] = 1 so when we get any subarray with sum = X, we can add prefSum[pref - X] = prefSum[0] = 1 to the answer.
Iterate over all the elements arr[i],
- Add arr[i] to prefix sum pref, pref += arr[i].
- Add the frequency of prefSums[pref - X] to the cnt.
- Increment the frequency of pref by 1.
Return the final answer as cnt.

Below is the implementation of the algorithm:

C++

#include <bits/stdc++.h>
#define ll long long int
using namespace std;

// function to find the number of subarrays with sum = X
ll solve(vector<ll>& arr, ll X, ll N)
{
    // Map to store the frequency of prefix sums
    map<ll, ll> prefSums;

    prefSums[0] = 1;
    ll pref = 0;
    ll cnt = 0;

    // Calculate the prefix sum at every index, and find the
    // count of subarrays with sum = pref - X
    for (int i = 0; i < N; ++i) {
        pref += arr[i];
        cnt += (prefSums[pref - X]);
        prefSums[pref] += 1;
    }
    return cnt;
}

int main()
{
    // Sample Input
    ll N = 5, X = 3;
    vector<ll> arr = { 1, 1, 1, -1, 1 };

    cout << solve(arr, X, N);
    return 0;
}

Java

import java.util.HashMap;
import java.util.Map;

public class Main {
    public static int GFG(int[] arr, int X, int N) {
        // Map to store the frequency of the prefix sums
        Map<Integer, Integer> prefSums = new HashMap<>();
        prefSums.put(0, 1);
        int pref = 0;
        int cnt = 0;

        // Calculate the prefix sum at every index and find the
        // count of the subarrays with sum = pref - X
        for (int i = 0; i < N; ++i) {
            pref += arr[i];
            cnt += prefSums.getOrDefault(pref - X, 0);
            prefSums.put(pref, prefSums.getOrDefault(pref, 0) + 1);
        }

        return cnt;
    }

    public static void main(String[] args) {
        // Sample Input
        int N = 5, X = 3;
        int[] arr = {1, 1, 1, -1, 1};
        System.out.println(GFG(arr, X, N));
    }
}

using System;
using System.Collections.Generic;

public class Solution
{
    public static int GFG(int[] arr, int X, int N)
    {
        // Dictionary to store the frequency of the prefix sums
        Dictionary<int, int> prefSums = new Dictionary<int, int>();
        prefSums.Add(0, 1);
        int pref = 0;
        int cnt = 0;

        // Calculate the prefix sum at every index and find the
        // count of the subarrays with sum = pref - X
        for (int i = 0; i < N; ++i)
        {
            pref += arr[i];
            cnt += prefSums.ContainsKey(pref - X) ? prefSums[pref - X] : 0;
            if (prefSums.ContainsKey(pref))
                prefSums[pref] = prefSums[pref] + 1;
            else
                prefSums.Add(pref, 1);
        }

        return cnt;
    }

    public static void Main(string[] args)
    {
        // Sample Input
        int N = 5, X = 3;
        int[] arr = { 1, 1, 1, -1, 1 };
        Console.WriteLine(GFG(arr, X, N));
    }
}

Javascript

function GFG(arr, X, N) {
    // Map to store the frequency of the prefix sums
    let prefSums = new Map();
    prefSums.set(0, 1);
    let pref = 0;
    let cnt = 0;
    // Calculate the prefix sum at every index and find the
    // count of the subarrays with sum = pref - X
    for (let i = 0; i < N; ++i) {
        pref += arr[i];
        cnt += (prefSums.get(pref - X) || 0);
        prefSums.set(pref, (prefSums.get(pref) || 0) + 1);
    }
    return cnt;
}
// Sample Input
let N = 5, X = 3;
let arr = [1, 1, 1, -1, 1];
console.log(GFG(arr, X, N));

Python3

# function to find the number of subarrays with sum = X
def solve(arr, X, N):
    # Dictionary to store the frequency of prefix sums
    pref_sums = {0: 1}
    pref = 0
    cnt = 0
    
    # Calculate the prefix sum at every index, and find the
    # count of subarrays with sum = pref - X
    for i in range(N):
        pref += arr[i]
        cnt += pref_sums.get(pref - X, 0)
        pref_sums[pref] = pref_sums.get(pref, 0) + 1
    
    return cnt
  
# Driver code  

if __name__ == "__main__":
    
    N = 5
    X = 3
    arr = [1, 1, 1, -1, 1]

    print(solve(arr, X, N))

Output

Time Complexity: O(N * logN), where N is the size of array arr[].
Auxiliary Space: O(N)

Article Tags :

Competitive Programming

DSA

Map

CSES Problems

prefix-sum