Java Concepts and DSA

Counting Subarrays with Sum Equal to K: A Prefix Sum Approach in Java

Anurag Srivastava's photo
Anurag Srivastava
·

4 min read

Finding subarrays with a specific sum is a classic problem in computer science. In this blog post, we'll explore an efficient solution using the concept of prefix sums and a HashMap in Java.

Understanding the Problem

Given an array of integers num and an integer k, we need to find the number of contiguous subarrays whose sum is equal to k.

The Prefix Sum Concept

The prefix sum of an array at index i is the sum of all elements from index 0 to i. Prefix sums allow us to quickly calculate the sum of any subarray in constant time.

The Algorithm

Our Java code utilizes the prefix sum concept along with a HashMap to efficiently count the subarrays with sum k:

  1. Initialize Variables:

    • preSum: Stores the running prefix sum.

    • hm: A HashMap to store prefix sums and their frequencies.

    • cnt: Counts the number of subarrays with sum k.

    • Initialize hm with (0, 1) because an empty subarray has a sum of 0.

  2. Iterate Through the Array:

    • For each element num[i] in the array:

      • Update preSum by adding num[i].

      • Calculate remove = preSum - k. If hm contains remove, it means there's a prefix sum that, when subtracted from the current preSum, equals k. This indicates a subarray with sum k.

      • Increment cnt by the frequency of remove in hm.

      • Update the frequency of the current preSum in hm.

  3. Return the Count:

    • Return cnt, which represents the number of subarrays with sum k.

Java Code Implementation

Java

package Arrays;

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

public class CountSubArraySumK {

    public static int countSubArray(int num[], int k) {
        int preSum = 0;
        Map<Integer, Integer> hm = new HashMap<Integer, Integer>();
        hm.put(0, 1);
        int cnt = 0;

        for (int i = 0; i < num.length; i++) {
            preSum += num[i];
            int remove = preSum - k;
            if (hm.containsKey(remove)) {
                cnt += hm.get(remove);
            }
            if (hm.containsKey(preSum)) {
                hm.put(preSum, hm.get(preSum) + 1);
            } else {
                hm.put(preSum, 1);
            }
        }
        return cnt;
    }

    public static void main(String[] args) {
        int num[] = {3, 1, 2, 4, -4, -2};
        System.out.println(countSubArray(num, 6));

        int num2[] = {1, 1, 1};
        System.out.println(countSubArray(num2, 2));

        int num3[] = {1, 2, 3};
        System.out.println(countSubArray(num3, 3));

        int num4[] = {1, -1, 0};
        System.out.println(countSubArray(num4, 0));
    }
}

Step-by-Step Examples

Example 1: num = {3, 1, 2, 4, -4, -2}, k = 6

  1. Initialization:

    • preSum = 0, hm = {0: 1}, cnt = 0

  2. Iteration:

    • i = 0, num[0] = 3, preSum = 3, remove = -3, hm = {0: 1, 3: 1}

    • i = 1, num[1] = 1, preSum = 4, remove = -2, hm = {0: 1, 3: 1, 4: 1}

    • i = 2, num[2] = 2, preSum = 6, remove = 0, cnt = 1, hm = {0: 1, 3: 1, 4: 1, 6: 1}

    • i = 3, num[3] = 4, preSum = 10, remove = 4, cnt = 2, hm = {0: 1, 3: 1, 4: 2, 6: 1, 10: 1}

    • i = 4, num[4] = -4, preSum = 6, remove = 0, cnt = 3, hm = {0: 1, 3: 1, 4: 2, 6: 2, 10: 1}

    • i = 5, num[5] = -2, preSum = 4, remove = -2, cnt = 3, hm = {0: 1, 3: 1, 4: 3, 6: 2, 10: 1}

  3. Result:

    • cnt = 3 (subarrays: [3, 1, 2], [2, 4], [4, -4, -2, 4])

Example 2: num = {1, 1, 1}, k = 2

  1. Initialization:

    • preSum = 0, hm = {0: 1}, cnt = 0

  2. Iteration:

    • i = 0, num[0] = 1, preSum = 1, remove = -1, hm = {0: 1, 1: 1}

    • i = 1, num[1] = 1, preSum = 2, remove = 0, cnt = 1, hm = {0: 1, 1: 1, 2: 1}

    • i = 2, num[2] = 1, preSum = 3, remove = 1, cnt = 2, hm = {0: 1, 1: 2, 2: 1, 3: 1}

  3. Result:

    • cnt = 2 (subarrays: [1, 1], [1, 1])

Time and Space Complexity

  • Time Complexity: O(n), where n is the length of the array, as we iterate through the array once.

  • Space Complexity: O(n), in the worst case, as the HashMap can store up to n distinct prefix sums.

This approach provides an efficient way to count subarrays with a specific sum using prefix sums and a HashMap. Happy coding!

#PrefixSum