Taro Logo

Subarray Sums Divisible by K

Medium
Amazon logo
Amazon
2 views
Topics:
ArraysDynamic Programming

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 answer is 7 because the following subarrays have 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 answer is 0.

Write an efficient algorithm to solve this problem. What is the time and space complexity of your solution?

Solution


Brute Force Solution

A naive approach would be to check every possible subarray and see if the sum is divisible by k. We can iterate through all possible start indices and end indices of subarrays and compute the sum. If the sum is divisible by k, we increment a counter.

def subarraysDivByK_brute_force(nums, k):
    count = 0
    for i in range(len(nums)):
        for j in range(i, len(nums)):
            subarray_sum = sum(nums[i:j+1])
            if subarray_sum % k == 0:
                count += 1
    return count

Time Complexity: O(n^2), where n is the length of the input array nums. This is because we have nested loops to iterate through all possible subarrays.

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

Optimal Solution Using Prefix Sum and Hash Map

We can optimize this using the concept of prefix sums and a hash map. The idea is that if prefix_sum[i] % k == prefix_sum[j] % k, then the subarray nums[i+1:j+1] has a sum divisible by k. We store the frequency of each remainder in a hash map.

def subarraysDivByK_optimal(nums, k):
    remainder_counts = {0: 1}
    prefix_sum = 0
    count = 0
    for num in nums:
        prefix_sum += num
        remainder = prefix_sum % k
        if remainder in remainder_counts:
            count += remainder_counts[remainder]
            remainder_counts[remainder] += 1
        else:
            remainder_counts[remainder] = 1
    return count

Time Complexity: O(n), where n is the length of the input array nums. We iterate through the array only once.

Space Complexity: O(k), where k is the divisor. In the worst case, the hash map remainder_counts can store up to k different remainders (0 to k-1).

Edge Cases

  • Empty Array: If the input array is empty, the number of subarrays with a sum divisible by k is 0. However, the problem statement says that 1 <= nums.length, so we don't need to explicitly check for this.
  • k = 1: If k is 1, then the sum of every subarray is divisible by k, so the answer is the number of non-empty subarrays, which is n * (n + 1) // 2.
  • Negative Numbers: The code handles negative numbers correctly because the modulo operator in Python handles negative numbers such that x % k will always return a value between -abs(k-1) and abs(k-1). We account for this by tracking the remainders which can be negative. The dictionary keys can store negative remainders.
  • Large input arrays and large K: the algorithm is efficient for large inputs and K, as it iterates through the array only once, and the size of the dictionary storing the remainders is limited by K, not by the length of the input array.