Subarray Sums Divisible by K

Subarray Sums Divisible by K

January 28, 2024
medium
Prefix Sum

Problem Statement #

Given an array of integers nums and an integer k, return the total number of continuous subarrays whose sum is divisible by k.

Example:

  1. Input: nums = [4,5,0,-2,-3,1], k = 5 Output: 7 Explanation: There are 7 subarrays with a sum divisible by 5: [5], [5, 0], [5, 0, -2, -3], [0], [0, -2, -3], [-2, -3], [4, 5, 0, -2, -3, 1].

Solution Approach #

The solution involves using a hashmap to store the frequency of prefix sum mod k and applying the prefix sum technique.

Algorithm Steps #

  1. Create a hashmap to store the frequency of prefix sums modulo k.
  2. Initialize the hashmap with a zero sum frequency of 1.
  3. Iterate through the array, calculating the running prefix sum modulo k.
  4. Add the frequency of the current prefix sum to the total count.
  5. Update the hashmap with the new prefix sum.
  6. Return the total count.

Code (Python) #

def subarraysDivByK(nums, k):
    count, prefix_sum, total = 0, 0, 0
    count_map = {0: 1}

    for num in nums:
        prefix_sum = (prefix_sum + num) % k
        total += count_map.get(prefix_sum, 0)
        count_map[prefix_sum] = count_map.get(prefix_sum, 0) + 1

    return total

# Test the function
print(subarraysDivByK([4,5,0,-2,-3,1], 5))  # Output: 7

Time Complexity #

O(n), where n is the number of elements in nums. The array is processed in a single pass.

Space Complexity #

O(min(n, k)), for the hashmap storing prefix sum frequencies.