[LeetCode 筆記] 560. Subarray Sum Equals K

leetcode subarray sum equals k java python algorithm

Chris 2023-10-09 20:00:31 ‧ 829 瀏覽

分享至

前言

　　這題學習目標是 Prefix Sums 前綴和的概念， Prefix Sums 通常用於需要頻繁查詢陣列中某一區間的元素和的情況，這裡目標是找到一個陣列中連續子陣列的合回傳最大值，時間複雜度估為 O(n^2)，這裡有 JAVA 和 Python 的寫法。

題目

Given an array of integers nums and an integer k, return the total number of subarrays whose sum equals to k.
A subarray is a contiguous non-empty sequence of elements within an array.

給定一個整數陣列 nums 和一個整數 k，回傳總合等於 k 值的連續子陣列的數量。
子陣列是陣列中連續的非空元素陣列。

題目連結：https://leetcode.com/problems/subarray-sum-equals-k/

題目限制

1 <= nums.length <= 2 * 104
1000 <= nums[i] <= 1000
107 <= k <= 107

範例

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

[1, 1] ⇒ 2 ，符合 k 值的大小
[1, 1] ⇒ 2 ，符合 k 值的大小
所以 nums 陣列可以有兩個子陣列，Output 回傳 2

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

[1, 2] ⇒ 3 ，符合 k 值的大小
[3] ⇒ 3 ，符合 k 值的大小
所以 nums 陣列可以有兩個子陣列，Output 回傳 2

思路&筆記

這題主要的思路是 Prefix Sums 為某一個區間和的算法

我們先把符合 nums 陣列其中一個連續子陣列，拿出來做說明回推

假如有一連續子陣列 prefix_sum = [2, 2, 3]

然後把這個子陣列的前綴和的終點 - 前綴和的起點，終點就是 3 ，起點就是 2

這裡概念上不是數學上的減去，而是空間上的減去

如果獲取的子陣列 prefix_sum 的前綴和相減等於 k 的話 count++

筆者用示意圖可能會比較清晰

// for example

nums = [1, 1, 1, 1, 2, 2, 3, 5, 6, 7, 8]
k = 7

prefix_sum = [2, 2, 3] = 7
prefix_sum_end = [1, 1, 1, 1, 2, 2, 3] = 11
prefix_sum_start = [1, 1, 1, 1] = 4
prefix_sum_end - sum_prefix_start => [1, 1, 1, 1, 2, 2, 3] - [1, 1, 1, 1] = 7

prefix_sum = k 成立，count++

JAVA 實現

public class Solution {
    public int subarraySum(int[] nums, int k) {

        int count = 0;
        int[] sum = new int[nums.length + 1]; // 設定前綴和的陣列
        sum[0] = 0; // 設定陣列第一個索引為 0

        // 把陣列中所有的元素加起來，組合成各個前綴和
        for (int i = 1; i <= nums.length; i++)
            sum[i] = sum[i - 1] + nums[i - 1]; // sum = [0, 1, 1+2, 1+2+3]
      
        for (int start = 0; start < sum.length; start++) {
            for (int end = start + 1; end < sum.length; end++) {
                if (sum[end] - sum[start] == k) // 前綴和的終點 - 前綴和的起點
                    count++;
            }
        }
      
        return count;
    }
}

Python 實現

使用了和 Java 一樣的邏輯執行，換成 Python 的寫法，不過超過了時間限制

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:

        count = 0
        sum = [0] * (len(nums)+1) # 設定前綴和的陣列
        sum[0] = 0 # 設定陣列第一個索引為 0

        # 把陣列中所有的元素加起來，組合成各個前綴和
        for i in range(len(nums)+1):
            sum[i] = sum[i-1] + nums[i-1] # sum = [0, 1, 1+2, 1+2+3]

        for start in range(len(sum)):
            for end in range(start+1, len(sum)):
                if sum[end] - sum[start] == k: # 前綴和的終點 - 前綴和的起點
                    count+=1
        return count

另一種寫法加入 python 的字典來演算，作用是幫助我們統計前綴和出現的次數

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:

        count = 0 
        prefix_sum = 0 # 當前的前綴和

        # 字典 d 的 key 是前綴和，value 是該前綴和出現的次數
        d = {0 : 1} # 初始將 0 的出現次數設為 1，表示前綴和為 0 的情況

        for num in nums:

            prefix_sum = prefix_sum + num # 更新當前的前綴和

            # 因為已經設定為 0 ，說明存在一個前綴和相減為 0 ，代表和 k 一樣
            if prefix_sum - k in d:
                count = count + d[prefix_sum - k] # count 次數：0 + 1

            # 如果不存在，表示這是首次遇到這個前綴和，將該前綴和作為 key，將其出現次數設為 1。
            if prefix_sum not in d:
                d[prefix_sum] = 1
            else:
                d[prefix_sum] = d[prefix_sum] + 1 # 代表已經有其他前綴和等於當前的前綴和，把相同的前綴和 + 1

        return count