前缀和技巧
前綴和技巧
先看一道題
 
一道簡單卻十分巧妙的算法問題:算出一共有幾個和為 k 的子數組。
那我把所有子數組都窮舉出來,算它們的和,看看誰的和等于 k 不就行了。關鍵是,如何快速得到某個子數組的和呢,比如說給你一個數組 nums,讓你實現一個接口 sum(i, j),這個接口要返回 nums[i…j] 的和,而且會被多次調用,你怎么實現這個接口呢?
因為接口要被多次調用,顯然不能每次都去遍歷 nums[i…j],有沒有一種快速的方法在 O(1)O(1)O(1) 時間內算出 nums[i…j] 呢?這就需要前綴和技巧了。
一、什么是前綴和
前綴和的思路是這樣的,對于一個給定的數組 nums,我們 額外開辟一個前綴和數組進行預處理:
int n = nums.length; // 前綴和數組 int[] preSum = new int[n + 1]; preSum[0] = 0; for (int i = 0; i < n; i++)preSum[i + 1] = preSum[i] + nums[i];這個前綴和數組 preSum 的含義也很好理解,preSum[i] 就是 nums[0..i-1] 的和。那么如果我們想求 nums[i..j] 的和,只需要一步操作 preSum[j+1]-preSum[i] 即可,而不需要重新去遍歷數組了。
回到開始的問題,我們想求有多少個子數組的和為 k,借助前綴和技巧很容易寫出一個解法:
int subarraySum(int[] nums, int k) {int n = nums.length;// 構造前綴和int[] sum = new int[n + 1];sum[0] = 0; for (int i = 0; i < n; i++)sum[i + 1] = sum[i] + nums[i];int ans = 0;// 窮舉所有子數組for (int i = 1; i <= n; i++)for (int j = 0; j < i; j++)// sum of nums[j..i-1]if (sum[i] - sum[j] == k)ans++;return ans; }這個解法的時間復雜度 O(N2)O(N^2)O(N2) 空間復雜度 O(N)O(N)O(N),并不是最優的解法。
二、優化解法
優化這段代碼:
for (int i = 1; i <= n; i++)for (int j = 0; j < i; j++)if (sum[i] - sum[j] == k)ans++;第二層 for 循環在干嘛呢?翻譯一下就是,在計算,有幾個 j 能夠使得 sum[i] 和 sum[j] 的差為 k。毎找到一個這樣的 j,就把結果加一。
我們可以把 if 語句里的條件判斷移項,這樣寫:
if (sum[j] == sum[i] - k)ans++;優化的思路是:我直接記錄下有幾個 sum[j] 和 sum[i] - k 相等,直接更新結果,就避免了內層的 for 循環。
我們先來分析一下, 如果區間[left, right]中數字之和等于k, 那么是不是會有sum(right) - sum(left) == k,這樣就變成了sum(right) - k == sum(left)
我們用一個字典當做累積和, 對于nums = [1,1,-2,2,4,-2,2], k=2來說,累積和[1,2,0,2,6,4,6],這里我們用字典dict記錄一下每個累積和出現的次數,考慮到了累積和剛好等于k的情況,將dict初始化為{0:1}
下面我們來遍歷一下nums,看看具體情況:
SUM = 0,SUM - k = -2 沒有出現在dict中, count = 0, dict = {0:1}(base case) SUM = 1, SUM - k = -1 沒有出現在dict中, count = 0, dict = {0:1, 1:1} SUM = 2, SUM - k = 0 出現在了dict中, count = 1, dict = {0:1,1:1,2:1} SUM = 0, SUM - k = -2 沒有出現在dict中, count = 1, dict = {0:2,1:1,2:1} SUM = 2, SUM - k = 0 出現在了dict中, count = 3, dict = {0:2,1:1,2:2} SUM = 6, SUM - k = 4 沒有出現在dict中, count = 3, dict = {0:2,1:1,2:2,6:1} SUM = 4, SUM - k = 2 出現在了dict中, count = 5, dict = {0:2,1:1,2:2,6:1,4:1} SUM = 6, SUM - k = 4 出現在了dict中, count = 6, dict = {0:2,1:1,2:2,6:2,4:1}所以count最后為6!!其實分析下來,我們發現dict的作用就是記錄sum(left)出現的次數這種方法只要遍歷一遍數組就行!!
我們可以用哈希表,在記錄前綴和的同時記錄該前綴和出現的次數。
int subarraySum(int[] nums, int k) {int n = nums.length;// map:前綴和 -> 該前綴和出現的次數HashMap<Integer, Integer> preSum = new HashMap<>();// base casepreSum.put(0, 1);int ans = 0, sum0_i = 0;for (int i = 0; i < n; i++) {sum0_i += nums[i];// 這是我們想找的前綴和 nums[0..j]int sum0_j = sum0_i - k;// 如果前面有這個前綴和,則直接更新答案if (preSum.containsKey(sum0_j))ans += preSum.get(sum0_j);// 把前綴和 nums[0..i] 加入并記錄出現次數preSum.put(sum0_i, preSum.getOrDefault(sum0_i, 0) + 1);}return ans; }比如說下面這個情況,需要前綴和 8 就能找到和為 k 的子數組了,之前的暴力解法需要遍歷數組去數有幾個 8,而優化解法借助哈希表可以直接得知有幾個前綴和為 8。
這樣,就把時間復雜度降到了 O(N)O(N)O(N)
三、總結
前綴和不難,卻很有用,主要用于處理數組區間的問題。
比如說,讓你統計班上同學考試成績在不同分數段的百分比,也可以利用前綴和技巧:
int[] scores; // 存儲著所有同學的分數 // 試卷滿分 150 分 int[] count = new int[150 + 1] // 記錄每個分數有幾個同學 for (int score : scores)count[score]++ // 構造前綴和 for (int i = 1; i < count.length; i++)count[i] = count[i] + count[i-1];這樣,給你任何一個分數段,你都能通過前綴和相減快速計算出這個分數段的人數,百分比也就很容易計算了
總結