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 n^{2} n2)，空间复杂度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；
如果是和能被整除，那么可以尝试记录每个数和目标数相除的余数。