Skip to content

Files

Latest commit

 

History

History
308 lines (260 loc) · 7.77 KB

560.subarray-sum-equals-k.md

File metadata and controls

308 lines (260 loc) · 7.77 KB

560. Subarray Sum Equals K

Clarification Questions

  • Is the input array sorted?

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

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
}
  • Time Complexity: O(n^3).
  • Space Complexity: O(1).

Optimized Brute Force

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^2).
  • Space Complexity: O(1).

Prefix + Hash

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[1] + ... + ...     + (              ... + nums[j])
       = nums[0] + nums[1] + ... + nums[i] + (nums[i + 1] + ... + nums[j])
       = pre[i]                            + (nums[i + 1] + ... + nums[j])

// We transform to
pre[i] + (nums[i + 1] + ... + nums[j]) = pre[j]
// Transform to
pre[i] + k                             = pre[j]

// Transform to
k = pre[j] - pre[i]

// 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, that is pre[i] exists. 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] with sum k, and we add the count of the subarray from hash table to the result.

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

Let's see an example:

// Example 1
[1, 2, 3, 4], k = 5
 i
prefixSum = 1
countMap.containsKey(1 - 5) = countMap.containsKey(-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) = countMap.containsKey(-2) = false
    countMap[3] = 1 {0:1, 1:1, 3:1}

[1, 2, 3, 4], k = 5
    ^^^^
       i
       prefixSum = 6
       countMap.containsKey(6 - 5) = countMap.containsKey(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 - 5) = countMap.containsKey(3) = true
          count += countMap[3] = 1
          countMap[10] = 1 {0:1, 1:1, 3:1, 6:1}

Implementation

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
}

Critical Points

There are some critical points 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 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}