Skip to content

560.subarray-sum-equals-k.md

Latest commit

 

History

History
295 lines (251 loc) · 7.63 KB

560.subarray-sum-equals-k.md

File metadata and controls

295 lines (251 loc) · 7.63 KB

560. Subarray Sum Equals K

(Optimized) Brute Force

fun subarraySum(nums: IntArray, k: Int): Int {
    var count = 0
    for (start in 0 until nums.size) {
        for (end in start until nums.size) {
            var sum = 0
            for (k in star..end) {
                sum += nums[k]
            }
            if (sum == k) count++
        }
    }
    return count
}

// We can sum up while iterating the element, because sum of `start..end` is the sum of `start..end - 1` + `nums[end]`.
fun subarraySum(nums: IntArray, k: Int): Int {
    var count = 0
    for (start in 0 until nums.size) {
        var sum = 0
        for (end in start until nums.size) {
            sum += nums[end]
            if (sum == k) count++
        }
    }
    return count
}
  • Time Complexity: O(n^3) and O(n^2).
  • Space Complexity: O(1).

Prefix Sum + Hash Table

We have prefix sum array pre[i], which is the sum of 0..i. And suppose we have another index j which i < j, then:

pre[j] = nums[0] + ...     + ... + ...     + (              ... + nums[j])
       = nums[0] + nums[1] + ... + nums[i] + (nums[i + 1] + ... + nums[j])
       = pre[i]                            + (nums[i + 1] + ... + nums[j])

// So we can transform the equation:
pre[j] = pre[i] + (nums[i + 1] + ... + nums[j])
pre[i] + k                             = pre[j]

// Transform to (the key formula)
pre[i] = pre[j] - k

To search pre[i] + k == pre[j], we can actually search if pre[j] - k exists (we had seen pre[i] before).

pre[j] - k = pre[i] indicates that sum of A[i+1 : j] == k

Let's see an example:

k = 3
i    =  0  1  2   3  4
nums = [2, 1, 3, -1, 1]
pre =  [2, 3, 6,  5, 6]
           i         j
        |--i
        |------------j
              |--k---|           

pre[i] = 2 + 1
pre[j] = 2 + 1  + 3 + -1 + 1
       = pre[i] + 3 + -1 + 1

pre[j] - pre[i] = 3 + -1 + 1
                = k

We can apply the same approach in problem 1. Two Sum, we maintain a hash table with prefix sum as key and the count of the prefix sum as value. Then we can iterate j to sum up the prefixSum (that is pre[j]), then check if pre[j] - k (pre[i]) exists in the hash table. If it exists, then we found a subarray A[i+1 : j] that can sum up to k, and we add the count of the subarray from hash table to the result.

Let's see an example:

// Example 1
[1, 2, 3, 4], k = 5
 i
prefixSum = 1
countMap.containsKey(1 - 5 = 4) = false
countMap[1] = 1 {0:1, 1:1} // Add the current prefixSum to the map

[1, 2, 3, 4], k = 5
    i
    prefixSum = 3
    countMap.containsKey(3 - 5 = -2) = false
    countMap[3] = 1 {0:1, 1:1, 3:1}

[1, 2, 3, 4], k = 5
    ^^^^
       i
       prefixSum = 6
       countMap.containsKey(6 - 5 = 1) = true
       count += countMap[1] = 1
       countMap[6] = 1 {0:1, 1:1, 3:1, 6:1}

// Example 2
[1, 2, 3, 4], k = 7
       ^^^^
          i
          prefixSum = 10
          countMap.containsKey(10 - 7 = 3) = true
          count += countMap[3] = 1
          countMap[10] = 1 {0:1, 1:1, 3:1, 6:1}
fun subarraySum(nums: IntArray, k: Int): Int {
    var count = 0
    var prefixSum = 0
    val countMap = HashMap<Int, Int>()
    // We don't set value here 
    // countMap[0] = 1
    for (i in 0 until nums.size) {
        prefixSum += nums[i]

        // We found the current `prefixSum == k`
        if (prefixSum == k) count++

        // We update the answer when we find the subarray 
        // with sum k by checking `prefixSum - k`
        if (countMap.containsKey(prefixSum - k)) {
            count += countMap[prefixSum - k]!!
        }
        // Increment the occurrence of current `prefixSum`
        countMap[prefixSum] = (countMap[prefixSum] ?: 0) + 1
    }
    return count
}

// Or equivalent
fun subarraySum(nums: IntArray, k: Int): Int {
    var count = 0
    var prefixSum = 0
    val countMap = HashMap<Int, Int>()

    // The CATCH: This is for the case for `prefixSum` == k or `prefixSum - k == 0`
    countMap[0] = 1

    for (i in 0 until nums.size) {
        prefixSum += nums[i]

        if (countMap.containsKey(prefixSum - k)) {
            count += countMap[prefixSum - k]!!
        }
        countMap[prefixSum] = (countMap[prefixSum] ?: 0) + 1
    }
    return count
}

Key Details

There are some critical points to pay attention to in our implementation and approach:

  • We don't maintain the prefixSum array because we can iterate once to calculate the prefixSum and put it in the hash table.
  • For the case prefixSum == k, we have to add the count of prefixSum == k to the result, so we initialize the prefixSum map with 0 as key and 1 as value.
k = 2
nums    = [1, 1, 2]
prefix  =  1  2  4
              * --> prefixSum - k = 0
  • If prefixSum - k exists, we don't increment the count by 1 because it leads to WA, there might be multiple prefixSum - k exists, so we sum up the count by the number of prefixSum - k exists. (Input [0, 0, 0], k = 0)
input = 0, 0, 0
prefix= 0, 0, 0
  • This problem can't be solved by sliding window, because the input might be negative and the sum won't be always increasing when expanding (monotonic). There is no good way to decide when to shrink the window when the number is negative (nums = [-1, -1, 1] and k = 0) or contains 0 (nums = [0, 1, 1, 0] and k = 2).
// Sliding Window, it doesn't work if nums contains negative number or 0, it only works if all numbers are positive.
fun subarraySum(nums: IntArray, k: Int): Int {
    if (nums.size == 1 && nums[0] != k) return 0
    var start = 0
    var end = 0
    var sum = 0
    var count = 0
    while (end < nums.size) {
        sum += nums[end]
        while (sum > k) {
            sum -= nums[start]
            start++
        }
        if (sum == k) {
            count++
        }
        end++
    }
    return count
}

Based on the characteristics of sliding window:

  • If wider window is valid, then narrower window is also valid.
k = 3
nums = [1, 1, 1]
        L     R     // valid
           L  R     // invalid
  • If narrow window is invalid, then wider window is invalid as well.
k = 2
nums = [1, 1, 1, -1]
        L     R     // invalid
           L  R     // valid
           L      R // invalid
  • Time Complexity: O(n).
  • Space Complexity: O(n).

Dry Run

Initial
count = 0
countMap[0] = 1

// ----------------
// Example 1
[5, 8], k = 5
 i
prefixSum = 5
countMap.containsKey(5 - 5) = countMap.containsKey(0) = true
count += countMap[0] = 1
countMap[5] = 1 {0:1, 5:1}

[5, 8], k = 5
    i
    prefixSum = 13
    countMap.containsKey(13 - 5) = countMap.containsKey(8) = false+
    countMap[13] = 1 {0:1, 5:1, 13:1}

// ----------------
// Example 2 
[5, 8], k = 8
 i
prefixSum = 5
countMap.containsKey(5 - 8) = countMap.containsKey(-3) = false
countMap[5] = 1 {0:1, 5:1}

[5, 8], k = 8
    i
    prefixSum = 13
    countMap.containsKey(13 - 8) = countMap.containsKey(5) = true
    countMap[13] = 1 {0:1, 5:1, 13:1}

// ----------------
// Example 3
[5, 8], k = 13
    i
    prefixSum = 13
    countMap.containsKey(13 - 8) = countMap.containsKey(5) = true
    countMap[13] = 1 {0:1, 5:1, 13:1}

Test Cases

Normal Cases

Input: nums = [1, 1, 1], k = 2
Output: 2

Input: nums = [1, 2, 3], k = 3
Output: 2

Edge / Corner Cases

  • All elements are 0:
Input: nums = [0, 0, 0], k = 0
Output: 6
  • Positive + Negative + Zero:
Input: nums = [0, 1, -1, 0], k = 0
Output: 5

Input: nums = [1, -1, 0, 1, -1], k = 0
Output: 7

Input: nums = [1, -1, 1, -1, 1], k = 0
  • All elements are equals to k:
Input: nums = [2, 2, 2], k = 2
Output: 3