LeetCode Solution, Medium, 560. Subarray Sum Equals K

Play this article

560. Subarray Sum Equals K

題目敘述

Given an array of integers nums and an integer k, return the total number of continuous subarrays whose sum equals to k.

Example 1:

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

Example 2:

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

Constraints:

  • 1 <= nums.length <= 2 * 10^4
  • -1000 <= nums[i] <= 1000
  • -10^7 <= k <= 10^7

Hint 1:

Will Brute force work here? Try to optimize it.

Hint 2:

Can we optimize it by using some extra space?

Hint 3:

What about storing sum frequencies in a hash table? Will it be useful?

Hint 4:

sum(i,j)=sum(0,j)-sum(0,i), where sum(i,j) represents the sum of all the elements from index i to j-1. Can we use this property to optimize it.

題目翻譯

這題是要找出在 nums 中的子陣列其元素是可以加總起來相當於 k 的組合。例如範例一, [1, 1, 1] 可以符合 k=2 的組合是 [1, 1, 1] 和 [1, 1, 1]。

解法解析

這題的解法是使用了 Hash Map 來處理,其 Time complexity 是最好的只有 O(n),雖然 Space complexity 只能到 O(n),但是如果要讓 Space complexity 降到 O(1) 的解法,卻會讓 Time complexity 漲到 O(n^2),實在不是一個划算的交易。

解題的概念是:

  1. 如果從索引 0 到索引 i 的總和和到索引 j 的總和相同,代表索引 i 到索引 j 的總和是 0。sum[i] + nums[i+1] + nums[i+...] = sum[j] 代表 nums[i+1] + nums[i+...] = 0
  2. 如果 sum[i] - sum[j] = k,代表 nums[i+1] + nums[i+...] = k

依照此概念,使用 hash map 去儲存所有的 sum[i] 出現過的次數,並且判斷 sum[i] - sum[j] = k 的總次數,就會是答案。

程式範例

Go
func subarraySum(nums []int, k int) int {
    count := 0
    sum := 0
    m := map[int]int{
        0: 1,
    }
    for _, num := range nums {
        sum += num
        if val, ok := m[sum-k]; ok {
            count += val
        }
        m[sum]++
    }

    return count
}
JavaScript
/**
 * @param {number[]} nums
 * @param {number} k
 * @return {number}
 */
var subarraySum = function (nums, k) {
    let count = 0,
        sum = 0;
    const map = new Map([[0, 1]]);
    for (const num of nums) {
        sum += num;
        count += map.get(sum - k) || 0;
        map.set(sum, (map.get(sum) || 0) + 1);
    }
    return count;
};
Kotlin
class Solution {
    fun subarraySum(nums: IntArray, k: Int): Int {
        var count = 0
        var sum = 0
        val map: HashMap<Int, Int> = hashMapOf(0 to 1)
        for (num in nums) {
            sum += num
            count += map.getOrDefault(sum - k, 0)
            map[sum] = map.getOrDefault(sum, 0) + 1
        }
        return count
    }
}
PHP
class Solution
{

    /**
     * @param Integer[] $nums
     * @param Integer $k
     * @return Integer
     */
    function subarraySum($nums, $k)
    {
        $count = 0;
        $sum = 0;
        $map = [0 => 1];
        foreach ($nums as $num) {
            $sum += $num;
            if (isset($map[$sum - $k])) {
                $count += $map[$sum - $k];
            }
            $map[$sum] = isset($map[$sum]) ? $map[$sum] + 1 : 1;
        }
        return $count;
    }
}
Python
class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        count, amount = 0, 0
        hash_map = {0: 1}
        for num in nums:
            amount += num
            count += hash_map.get(amount - k, 0)
            hash_map[amount] = hash_map.get(amount, 0) + 1
        return count
Rust
impl Solution {
    pub fn subarray_sum(nums: Vec<i32>, k: i32) -> i32 {
        let mut sum = 0;
        let mut count = 0;
        let mut map = std::collections::HashMap::new();
        map.insert(0, 1);
        for num in nums {
            sum += num;
            if let Some(v) = map.get(&(sum - k)) {
                count += v;
            }
            *map.entry(sum).or_insert(0) += 1;
        }
        return count;
    }
}
Swift
class Solution {
    func subarraySum(_ nums: [Int], _ k: Int) -> Int {
        var count = 0
        var sum = 0
        var map: [Int:Int] = [0:1]
        for num in nums {
            sum += num
            count += map[sum - k, default: 0]
            map[sum, default: 0] += 1
        }
        return count
    }
}

Did you find this article valuable?

Support 攻城獅 by becoming a sponsor. Any amount is appreciated!