Leetcode — 560. Subarray Sum Equals K (中文)

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

給予一整數陣列 nums 和一整數 k, 回傳可以找到幾個總和為k的連續子陣列(continuous subarray)

條件:

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

範例
Input: nums = [1,2,3], k = 3
Output: 2, 兩個連續子陣列[1,2]和[3]總和皆為3

想法:

我的第一個想法是直接用暴力解來解這題, 我們試著去組 從位置i為首的各種子陣列們, 然後看我們可不可以找到總和為k的子陣列, 這樣的方法會需要O(n²), 因為我們在每次位置i的時候, 都要嘗試看當子陣列的結尾位置在i~陣列最後一位的時候, 總和為何

以下為 1211ms 的程式碼, 只贏過 20% —

public int subarraySum(int[] nums, int k) {
        int count = 0;
        int n = nums.length;
        for(int i=0;i<n;i++)
        {
            int temp = 0;
            for(int j=i;j<n;j++)
            {
                temp += nums[j];
                if(temp==k)
                    count++;
            }
        }
        return count;
    }

雖然這樣的程式碼可以幫助我們通過Leetcode上的所有test case, 但我們是不是可以用一些空間來減少我們的時間複雜度

我們來試著想我們有沒有可能用O(n)的時間去解這個問題? 來先畫我們剛剛暴力解的過程, 看我們是否可以省去不必要的流程

👷假設我們的陣列 nums = [1, 2, 3, 4, 1]

⚠️ 為了待會解釋方便, 我們將用nums[i, j]的形式代表我們組成的子陣列, 其起始位置為i, 結束位置為j
比如說在這個範例中, nums[0,2]代表的就是子陣列[1,2,3]

我們把上面程式碼的過程畫下來如下圖, 寫下來所有子陣列nums[i, j]的總和, 圖中同一列就代表這些子陣列有相同起始位置, 但不同結尾位置 (而圖中這些橘色數字就是我們上面程式碼 利用temp變數所得到的總和們)

Example Graph

如果你仔細看上面的圖, 你可以發現除了第一行(也就是從最一開始位置組成的子陣列總和們)以外, 其他列的數字總和我們都可以靠從 nums[0, i] 減掉nums[0, j]得到, 且 0≤j<i (這兩者數字相減也代表我們拿到的就是子陣列nums[j+1, i]的總和)

用例子來說的話, 仔細看當我們子陣列結尾在位置3的總和們, 我們可以靠 用nums[0, 3]的子陣列 — [1,2,3,4] 的總和 減去前面的任一子陣列nums[0, j]如[1], [1,2], 或[1,2,3] 的總和, 可以得到總和9, 7, 跟4
也就是說, 這些總和代表的就是子陣列[2,3,4], [3,4],跟[4]的總和們

程式邏輯:

因此我們可以依照上面暴力解法的程式碼進行些微修改, 我們仍然需要剛剛用過的temp變數來儲存從0開始的子陣列總和(也就是上面圖中的第一列的數字們), 但這次我們不再用另一個迴圈來幫我們找新的總和們了, 我們將會用一個新變數來儲存我們目前見過的所有nums[0, j]的總和 (也就是我們在位置i的時候, 就會知道所有nums[0,j]的總和, j=0~i-1)

我決定用一個Set來儲存我們遇過的總和們, 因為在Java中, 用Set來搜尋特定數字的時間複雜度最低
我們需要用一個迴圈來跑在nums陣列中的各個位置i, 每次都必須做三件事情:

1. 把現在位置的數字nums[i] 累加進temp變數中 (因此temp會永遠代表nums[0, i]的總和)
2. 檢查現在的nums[0, i]減去先前遇過的所有總和nums[0, j] 是否可以得到我們目標總和k
3. 把現在剛加的nums[0, i]總和放進set變數

另外, 我們也必須把0先放進set中, 以免之後遇到情況像是nums[0, i]的總和已經是k了, 代表nums[0, i]可以自成總和k

依照上述邏輯就可以寫成下列程式碼:

public int subarraySum(int[] nums, int k) {
        int count = 0;
        int n = nums.length;
        Set<Integer> sums = new HashSet<>();
        sums.add(0);
        int temp = 0;
        for(int i=0;i<n;i++)
        {
            temp+= nums[i];
            if(sums.contains(temp-k))
                count++;
            sums.add(temp);
        }
        return count;
    }

但這個程式碼會在input nums = [-1,1,0,0], k = 0的時候失敗
很明顯的, 原因出在我沒有考慮到這兩個不同子陣列nums[0, j]跟nums[0, j’]可能會有相同的總和
因此, 我們不能僅僅用Set來儲存所有遇過的nums[0, i]總和, 我們將會用Java Map來儲存各遇過的總和,以及遇過這個總和的總次數

因此我們可以得到下面 19ms 的程式碼, 時間複雜度為 O(n) —

public int subarraySum(int[] nums, int k) {
        int count = 0;
        int n = nums.length;
        Map<Integer,Integer> map = new HashMap<>();//sum - count
        map.put(0, 1);
        int temp = 0;
        for(int i=0;i<n;i++)
        {
            temp+= nums[i];
            count += map.getOrDefault(temp-k,0);
            map.put(temp, map.getOrDefault(temp,0)+1);
        }
        return count;
    }