Taro Logo

Subarray Sums Divisible by K

Medium
Google logo
Google
1 view
Topics:
Arrays

You are given an integer array nums and an integer k. Your task is to find the number of non-empty subarrays that have a sum divisible by k. A subarray is defined as a contiguous part of an array.

For example:

  • nums = [4, 5, 0, -2, -3, 1], k = 5. The output should be 7 because the following subarrays have a sum divisible by 5: [4, 5, 0, -2, -3, 1], [5], [5, 0], [5, 0, -2, -3], [0], [0, -2, -3], [-2, -3].
  • nums = [5], k = 9. The output should be 0 because no subarray has a sum divisible by 9.

Could you please provide an algorithm to solve this problem efficiently, considering the following constraints?

  • 1 <= nums.length <= 3 * 10^4
  • -10^4 <= nums[i] <= 10^4
  • 2 <= k <= 10^4

Solution


Brute Force Solution

A naive approach would be to iterate through all possible subarrays and check if their sum is divisible by k. This involves two nested loops to define the start and end indices of the subarrays.

Algorithm:

  1. Iterate through all possible start indices i from 0 to n-1.
  2. For each start index i, iterate through all possible end indices j from i to n-1.
  3. Calculate the sum of the subarray from i to j.
  4. If the sum is divisible by k, increment the count.
  5. Return the final count.
def subarraysDivByK_brute_force(nums, k):
    count = 0
    n = len(nums)
    for i in range(n):
        for j in range(i, n):
            subarray_sum = sum(nums[i:j+1])
            if subarray_sum % k == 0:
                count += 1
    return count

Time Complexity: O(n^2) due to the nested loops.

Space Complexity: O(1) as we are only using a constant amount of extra space.

Optimal Solution

We can optimize this using the concept of prefix sums and remainders. The key idea is that if two prefix sums have the same remainder when divided by k, then the subarray between those two indices has a sum divisible by k.

Algorithm:

  1. Initialize a prefix sum prefix_sum to 0.
  2. Initialize a remainder_count dictionary (or array) to store the count of each remainder.
  3. Initialize the remainder_count[0] to 1, as an empty subarray has a sum of 0, which is divisible by any k.
  4. Iterate through the nums array.
  5. Update the prefix_sum by adding the current element.
  6. Calculate the remainder of the prefix_sum when divided by k. Make sure the remainder is non-negative.
  7. Increment the count by the value stored in the remainder_count for the current remainder.
  8. Update the remainder_count for the current remainder.
  9. Return the final count.
def subarraysDivByK(nums, k):
    remainder_count = {0: 1}
    prefix_sum = 0
    count = 0
    for num in nums:
        prefix_sum += num
        remainder = prefix_sum % k
        if remainder < 0:
            remainder += k  # Ensure remainder is non-negative
        if remainder in remainder_count:
            count += remainder_count[remainder]
            remainder_count[remainder] += 1
        else:
            remainder_count[remainder] = 1
    return count

Time Complexity: O(n) because we iterate through the array only once.

Space Complexity: O(k) because, in the worst case, the remainder_count dictionary might store k different remainders.

Edge Cases

  • Empty array: The algorithm should handle an empty input array gracefully, returning 0.
  • Negative numbers: The algorithm should correctly handle negative numbers in the input array.
  • k = 1: If k is 1, every subarray is divisible by k, so you would return n * (n + 1) / 2
  • Zero in array: The optimal solution handles zero in the array correctly.