LeetCode技巧篇（一）prefix sum 前綴和

介紹

前綴和（prefix sum）是算法題中比較實用的一種技巧，當算法題的背景是整數型數組且出現 “子數組和” 或者 “連續的子數組” 既可以考慮使用前綴和來求解會得到不錯的效果。
假設給定的數組A各個元素分別爲：

那麼我們可以得到一個前綴和數組B，通過累加A[0:i-1]得到B[i]:

在實現上，可以直接在數組B的基礎上累加即可，不需要遍歷一遍A。

實例

現在看一道簡單的應用，LeetCode 560. Subarray Sum Equals K。題目很簡單，找到連續子數組和爲K的子數組個數。傳統方法是兩次循環，時間複雜度O($n^{2}$)，空間複雜度O(1)。

    public int subarraySum(int[] nums, int k) {
        int count = 0;
        for (int start = 0; start < nums.length; start++) {
            int sum=0;
            for (int end = start; end < nums.length; end++) {
                sum+=nums[end];
                if (sum == k)
                    count++;
            }
        }
        return count;
    }

但是這裏如果使用前綴和搭配HashMap可以實現O(n)的時間複雜度。
首先，我們需要得到前綴和數組。
然後遍歷前綴和數組，在HashMap找到是否有元素等於sum[i]-k；如果存在，則count+=map.get(sum[i]-k)。
無論存在與否都將當前元素放進map中，並且更新該值的個數。 代碼如下：

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

再看一題換湯不換藥的題目，LeetCode 1343. Number of Sub-arrays of Size K and Average Greater than or Equal to Threshold。找到長度爲K且平均值不小於閾值的子數組的個數。這裏規定了子數組的長度那就更簡單了：
首先，得到前綴和數組；
然後遍歷前綴和數組，看間隔爲k的數組的平均值是否大於等於閾值。 代碼如下：

    public int numOfSubarrays(int[] arr, int k, int threshold) {
        int count =0;
        threshold = threshold * k;
        int[] temp = new int[arr.length+1];
        for(int i=1;i<=arr.length;i++){
            temp[i] = temp[i-1] +arr[i-1];
        }
        for(int i=1;i<=temp.length-k;i++){
           if(temp[i+k-1] - temp[i-1]>=threshold) count++;
        }
        return count;
    }

再看一題稍微有點技巧的，LeetCode 974 Subarray Sums Divisible by K。
這題乍一看和前面1343有點像，但是難度明顯更大——因爲現在沒有規定子數組的長度了，所以這裏要轉換思路。既然要看子數組和能不能被K整除，那麼我們把前綴和數組的每一個元素都對K取餘，那麼要想子數組能被K整除，只有一種可能，那就是餘數相同，因爲這裏要做減法；如果要像LeetCode 1010. Pairs of Songs With Total Durations Divisible by 60一樣和能被60整除的話，那就是判斷餘數之和了。因此這一題可以看成是 前綴和和餘數保存結合起來的題目。

    public int subarraysDivByK(int[] A, int K) {
        int arr[] =new int[A.length+1];
        for(int i=1;i<=A.length;i++){
            arr[i] = arr[i-1] +A[i-1];
        }
        int[] count = new int[K];
        for (int x: arr)
        //這裏取餘可能會返回負值，因此再加一層
            count[(x % K + K) % K]++;
        int ans = 0;
        for (int v: count)
        //假設餘數爲3的個數有3個，那麼可能有3種情況使得子數組和能被K整除
            ans += v * (v - 1) / 2;
        return ans;
    }

總結

一句話，問題裏有subarray sum字樣，可以試試prefix sum；
如果是要匹配個數，可以搭配HashMap；
如果是和能被整除，那麼可以嘗試記錄每個數和目標數相除的餘數。