CODEWITHSUNDEEP
javaalgorithm
Learner
Learner

Reputation: 21445

Find the count of subarrays whose sum is equal to given value

I am trying to understand a program where the task is to find how many sub arrays are possible whose sum is equal to given value.

Here is the working code with O(N) complexity taken from link:

static int findSubarraySum(int arr[], int K) {
    Map<Integer, Integer> prevSum = new HashMap<>();
    int res = 0;
    int currsum = 0;
    for (int i = 0; i < arr.length; i++) {
        currsum += arr[i];
        if (currsum == K)
            res++;
        if (prevSum.containsKey(currsum - K))
            res += prevSum.get(currsum - K);
        Integer count = prevSum.get(currsum);
        if (count == null)
            prevSum.put(currsum, 1);
        else
            prevSum.put(currsum, count + 1);
    }

    return res;
}

public static void main(String[] args) {
    int arr[] = {10, 2, -2, -20, 10};
    int sum = -10;
    int n = arr.length;
    System.out.println(findSubarraySum(arr, sum));
}

I am not able to understand the logic behind below code:

            if (prevSum.containsKey(currsum - K))   
                res += prevSum.get(currsum - K);

How the HashMap in above code is helpful in getting correct result.

Upvotes: 0

Views: 200

Answers (1)

Nico Schertler
Nico Schertler

Reputation: 32667

prevSum[i] contains the number of contiguous subarrays starting with the first element that sum to i. E.g., for the array 3, 4, 2, -4, 2, the entries would look as follows:

prevSum[3] = 1  // { { 3 } }
prevSum[5] = 1  // { { 3, 4, 2, -4 } }
prevSum[7] = 2  // { { 3, 4 }, { 3, 4, 2, -4, 2} }
prevSum[9] = 1  // { { 3, 4, 2 } }

If you have a current prefix sum currsum of 10 and you are looking for subarrays that sum to K=3, you want to remove subarrays that start with the first element and sum to 7 to get an appropriate subarray that sums to 3. Example:

3, 4, 2, -4, 2, 1, 2, -5, 2
                   ^
|------------------|
              currsum = 10

|--|  -> sums to 7
     |-------------| -> sums to 3

|------------| -> sums to 7
                |--| -> sums to 3

Hence, at this location, you have found two possible subarrays that sum to K. Therefore, res += prevSum.get(currsum - K).

Upvotes: 1

Related Questions