Taro Logo

Subarray Sums Divisible by K

Medium
Meta logo
Meta
0 views
Topics:
Arrays

Given an integer array nums and an integer k, return the number of non-empty subarrays that have a sum divisible by k.

A subarray is 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 k = 5:

  • [4, 5, 0, -2, -3, 1] (sum is 5)
  • [5] (sum is 5)
  • [5, 0] (sum is 5)
  • [5, 0, -2, -3] (sum is 0)
  • [0] (sum is 0)
  • [0, -2, -3] (sum is -5)
  • [-2, -3] (sum is -5)

Another example:

nums = [5], k = 9

The output should be 0.

How would you approach this problem? Consider the time and space complexity of your solution. Can you provide both a brute force solution and a more optimized solution?

Solution


Brute Force Solution

The most straightforward approach is to check every possible subarray and count the ones whose sum is divisible by k.

  1. Iterate through all possible start indices of the subarray.
  2. For each start index, iterate through all possible end indices.
  3. Calculate the sum of the subarray.
  4. Check if the sum is divisible by k. If it is, increment the count.
def subarraysDivByK_brute_force(nums, k):
    count = 0
    n = len(nums)
    for i in range(n):
        for j in range(i, n):
            sub_array_sum = sum(nums[i:j+1])
            if sub_array_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 use the concept of prefix sums and the property that (a - b) % k == 0 if and only if a % k == b % k. We store the remainders in a hash map or an array.

  1. Initialize a hash map (or array) to store the count of each remainder when dividing by k. Initialize the count of remainder 0 to 1, as an empty subarray has a sum of 0, which is divisible by any k.
  2. Iterate through the nums array, keeping track of the cumulative sum.
  3. For each element, calculate the remainder of the cumulative sum when divided by k.
  4. If the remainder is negative, add k to it to make it non-negative (since negative remainders in Python can be different from other languages).
  5. Check if the remainder is already present in the hash map. If it is, add its count to the total count.
  6. Increment the count of the current remainder in the hash map.
def subarraysDivByK(nums, k):
    count = 0
    prefix_sums = {0: 1}
    curr_sum = 0

    for num in nums:
        curr_sum = (curr_sum + num) % k
        if curr_sum < 0:
            curr_sum += k

        if curr_sum in prefix_sums:
            count += prefix_sums[curr_sum]
            prefix_sums[curr_sum] += 1
        else:
            prefix_sums[curr_sum] = 1

    return count

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

Space Complexity: O(k) in the worst case, as the hash map can store at most k different remainders (0 to k-1).

Edge Cases

  • Empty array: The problem statement specifies non-empty subarrays, so we don't need to handle an empty input array explicitly.
  • Negative numbers: The code handles negative numbers by adding k to the remainder to make it non-negative.
  • Large values: The problem constraints limit the size of the input values and k, so integer overflow is not a major concern, but it's always good to keep in mind for extremely large inputs.
  • k = 1: If k = 1, all subarrays will have a sum divisible by k, so the result will be n*(n+1) / 2.

Example

nums = [4, 5, 0, -2, -3, 1], k = 5

  1. curr_sum = 4, remainder = 4. prefix_sums = {0: 1, 4: 1}
  2. curr_sum = 9, remainder = 4. count += 1, prefix_sums = {0: 1, 4: 2}
  3. curr_sum = 9, remainder = 4. count += 1, prefix_sums = {0: 1, 4: 2}
  4. curr_sum = 9, remainder = 4. count += 1, prefix_sums = {0: 1, 4: 2}
  5. curr_sum = 9, remainder = 4. count += 1, prefix_sums = {0: 1, 4: 2}
  6. curr_sum = 9, remainder = 4. count += 1, prefix_sums = {0: 1, 4: 2}