Taro Logo

Subarray Sums Divisible by K

Medium
a month ago

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.

Example 1:

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

Example 2:

Input: nums = [5], k = 9
Output: 0

Constraints:

  • 1 <= nums.length <= 3 * 10^4
  • -10^4 <= nums[i] <= 10^4
  • 2 <= k <= 10^4
Sample Answer
## Problem: Subarray Sums Divisible by K

Given an integer array `nums` and an integer `k`, the goal is to find the number of non-empty subarrays whose sum is divisible by `k`.

### Naive Approach

A brute-force solution involves iterating through all possible subarrays and checking if their sum is divisible by `k`. This approach has a time complexity of O(n^2).

```python
def subarraysDivByK_naive(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

Optimal Approach

An optimized approach uses the concept of prefix sums and remainders. We can keep track of the prefix sums modulo k. If two prefix sums have the same remainder when divided by k, it means the subarray between those two indices has a sum divisible by k.

def subarraysDivByK(nums, k):
    remainder_count = {0: 1}  # Initialize with remainder 0 appearing once (empty subarray)
    prefix_sum = 0
    count = 0

    for num in nums:
        prefix_sum = (prefix_sum + num) % k
        if prefix_sum in remainder_count:
            count += remainder_count[prefix_sum]
            remainder_count[prefix_sum] += 1
        else:
            remainder_count[prefix_sum] = 1

    return count

Example

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

  1. Initialize remainder_count = {0: 1}, prefix_sum = 0, count = 0
  2. Iterate through nums:
    • num = 4: prefix_sum = (0 + 4) % 5 = 4. remainder_count = {0: 1, 4: 1}
    • num = 5: prefix_sum = (4 + 5) % 5 = 4. count += remainder_count[4] = 1, remainder_count = {0: 1, 4: 2}
    • num = 0: prefix_sum = (4 + 0) % 5 = 4. count += remainder_count[4] = 1 + 2 = 3, remainder_count = {0: 1, 4: 3}
    • num = -2: prefix_sum = (4 - 2) % 5 = 2. remainder_count = {0: 1, 4: 3, 2: 1}
    • num = -3: prefix_sum = (2 - 3) % 5 = -1 % 5 = 4. count += remainder_count[4] = 3 + 3 = 6, remainder_count = {0: 1, 4: 4, 2: 1}
    • num = 1: prefix_sum = (4 + 1) % 5 = 0. count += remainder_count[0] = 6 + 1 = 7, remainder_count = {0: 2, 4: 4, 2: 1}
  3. Return count = 7

Big O Runtime Analysis

The optimal approach iterates through the array once. The time complexity is O(n) because each element is visited only once. The operations inside the loop (modulo, dictionary lookups, and updates) take constant time.

The naive approach involves nested loops, giving it a time complexity of O(n^2).

Big O Space Usage Analysis

The optimal approach uses a dictionary (remainder_count) to store the counts of remainders. In the worst case, all remainders are unique, so the dictionary can store up to k entries. Thus, the space complexity is O(k). In the worst case, k could be equal to n, making the space complexity O(n).

The naive approach has a space complexity of O(1) as it only uses a few constant space variables.

Edge Cases

  1. Empty array: If the input array is empty, the number of subarrays with a sum divisible by k is 0. The given code handles this case correctly.
  2. k = 1: If k is 1, any subarray's sum will be divisible by k. The number of subarrays is n(n+1)/2. The code correctly handles this case by counting all the possible subarrays.
  3. Large k: If k is very large compared to the numbers in the array, the remainders will likely be unique, increasing the space usage of the dictionary. The space complexity is still bounded by O(k).
  4. Negative numbers: The code correctly handles negative numbers using the modulo operator.
  5. Zero sum subarrays: If there are zero sum subarrays, these will be counted correctly, as a sum of 0 is divisible by any k.