Leetcode__1508. Range Sum of Sorted Subarray Sums

Leetcode__1508. Range Sum of Sorted Subarray Sums

这一题是一道十分典型的二分查找与双指针的使用，题目的大意是，给出一个数组 (nums) , 我们可以通过该数组得到子数组, 子数组的大小从 1 到 (n), 因此有 (frac{n(n+1)}{2}) 个不同的子数组, 然后我们要求这些子数组的和, 然后对这些子数组的和进行排序, 求排序后的数组的区间和, 那么这道题的结果就是一些子数组和 的和, 可能有一点拗口, 我们会用例子说明.

我们定义原数组为 (nums), 他的 (frac{n(n+1)}{2}) 个子数组的和 构成新的有序数组为 (SUM_{sub}) , 例如 (nums = [1,2,3,4]), 那么我们可以计算得到 (SUM_{sub} = [1, 2, 3, 3, 4, 5, 6, 7, 9, 10]) , 然后我们要在 (SUM_{sub}) 上求区间和, 传统的方式就是求得 (SUM_{sub}), 然后在区间 ([left: right]) 上求和.

问题思路转换

我们要将定义重新捋一下, (SUM_{sub}) 是子数组的和, 并且 (SUM_{sub}) 是有序的, 那么 (SUM_{sub}) 的前缀和(Prefix_{subsum}[k] =sum_{i=0}^k SUM_{sub}[i])就表示在子数组的和当中的最小的 (k) 个元素之和, 我们要求的是 (SUM_{sub}[left]) 到 (SUM_{sub}[right]) 这部分的和, 也就等于 (Prefix_{subsum}[right] - Prefix_{subsum}[left-1]),

现在思考一个问题， 如果找到了第 (k) 大的子数组和 (SUM_k) , 能否在(O(n))时间内找到小于 (SUM_k) 的子数组的和.

方法是: 对于原数组 (nums), 遍历下标 (0-i), 我们计算最大的 (j), 使得 (nums) 在区间([i:j]) 上的和 小于 (SUM_k), 那么在区间 ([i:j]) 内的所有以 (i) 为起始点的子数组 的子数组和 都小于 (SUM_k). (注意,这里下标 (j) 是包括, 官方题解是不包括, 当然, 这是个小问题)

例如 (nums = [1,2,3,4]), 如果我们知道了第 (8) 大的子数组和是 (7), 那么我们可以遍历 ([1,2,3,4]) 的下标, 从下标为 (0) 开始, 下标为(0) 的 (j) 是 (2), 下标 (1) 对应的(j) 是 (2), 那么就得到子数组和小于(6) 的子数组集合为 ({[1],[1,2],[1,2,3],[2],[2,3],[3],[4] }), 然后我们求这些子数组的和 的和就是 (Prefix_{subsum}[7]). 注意, 这里是 (Prefix_{subsum}[7]) 而不是 (8) 因为我们求的是严格小于 (SUM_8) 的子数组和, 上面的集合中, 刚好就是 (7) 个元素.

优化一下计算的方式

上述过程, 在求得 (SUM_k) 之后, 我们可以计算得到子数组和小于 (SUM_k) 的子数组集合的区间是: ({[0,0_j], [1 : 1_j], ..[m : m_j]}), 对于其中一个区间 ([m: m_j]) 来说, 一共有 (m_j - m +1) 个子数组, 现在我们需要优化计算这 (m_j - m +1) 个子数组和的和, 例如, 在上面的例子中, 对于 (m==0) 的时候, 子数组就是 ([1]) 和 ([1,2]), 子数组和 的和就是 (4).

计算方式就是引入前缀和, 定义 (Prefix_{sum}) 为 (nums) 的前缀和, (Prefix\_Prefix_{sum}) 是 (Prefix_{sum}) 的前缀和, 也就是 (nums) 前缀和的前缀和, 那么区间 ([i;j]) 上的所有以 (i) 为起始下标的 (j-i+1) 个子数组和的和为:

[Prefix\_Prefix_{sum}[j]−Prefix\_Prefix_{sum}[i-1]−Prefix_{sum}[i-1]×(j−i+1) ]

这里的(Prefix\_Prefix_{sum}[j]) 表示的是 (Prefix_{sum}) 的前缀和, (包括 (j)) , (Prefix\_Prefix_{sum}[j]−Prefix\_Prefix_{sum}[i-1]) 表示的是 (sum_{k = i}^{j} Prefix_{sum}[k]) , 但是这并不等于 ([i: j]) 上子数组和的和, 因为(Prefix_{sum}) 是前缀和, 那么会把 (nums[0]+ nums[1]...+nums[i-1]) 也计算在内, 计算的次数就是 (j-i+1) 次, 所以我们要减去 (Prefix_{sum}[i-1]).

例如

(nums = [1,2,3,4]), 那么

(Prefix_{sum}) 为 ([1,3,6,10]),

(Prefix\_Prefix_{sum}) 为 ([1,4,10,20]),

对于区间 ([2: 3]), (Prefix\_Prefix_{sum}[3]−Prefix\_Prefix_{sum}[1]) 等于 (9), 表示 (Prefix_{sum}) 数组中 (3, 6) 的和, 那么对于原数组来说 (Prefix_{sum}[0]) 就被多加进去了, 所以要删去, 删去的次数就是两次.

特殊情况((i== 0))

当 (i== 0) 时, (Prefix\_Prefix_{sum}[j]) 就是区间 ([0,0_j]) 上的子数组和的和, 不需要减去后半部分.

返回结果

我们已经计算了区间 ([i: j]) 上的(j-i+1) 个子数组和 的和, 回头看, 我们已知子数组和小于 (SUM_k) 的子数组集合的区间是: ({[0,0_j], [1 : 1_j], ..[m : m_j]}), 对于每个区间, 我们都使用上述的计算方式, 计算每个区间内的子数组和 的和. 我们就得到了整个 (nums) 数组中, 所有的 子数组和 小于(SUM_k) 的这些子数组和 的和. 换句话说, 就是 (sum_{i= 0}^{k_{low}}SUM_i) , 注意, 我们的计算方式并不是计算每个 (SUM_i), 而是用子区间的方式计算所有 (SUM_i) 的和.

还原假设

我们刚才的计算步骤都基于一个假设, 那就是 如果找到了第 (k) 大的子数组和 (SUM_k) , 现在我要做的就是如何找到 (SUM_k), 因此, 这一题的思路就是,

1. 使用二分法找到 (SUM_{right}) 和 (SUM_{left-1}), 注意, 这里的 (SUM_k) 是第 (k) 大, 而不是前 (k) 大的和,
2. 根据 (SUM_k), 计算前 (k) 大的和. 也就是计算前 (right) 大的和, 与 前 (left-1) 大的和.

二分法求 (SUM_k)

(SUM_k)的计算方式本质上和 计算小于 (SUM_k) 的区间的计算方式相同, 只不过, 这里我们是需要计算 (k), 而不是计算子数组和的和. 也就是说, 对于 (Middle), 我们要判断, 他是不是第 (k) 大, 那么我们计算 子数组和 小于 (Middle) 的子数组的个数 (Middle_{count}) 就可以了, 如果这个(Middle_{count}) 大于 (k), 那么我们就要缩小 (Middle), 如果 (Middle_{count}) 小于 (k), 我们就要增大 (Middle), 缩小与增大的方式就使用 二分法.

具体怎么求 (Middle) 对应的 (Middle_{count}) 呢, 本质上也是计算区间集合, 同样, 我们遍历 (i) 从 (0) 到 ((n-1)), 以(i) 为下标, 计算最大的 (j), 使得 (sum_{k=i}^{j} nums[k] <= Middle) ,也就是计算不超过 (Middle) 的子数组和对应的 (j). 注意,这里是小于等于, 前面讲的是严格小于. 那么从 (i) 到 (j) 一共有 (j-i+1) 个子数组. 遍历之后, 我们会得到子数组和小于 (Middle) 的子数组集合的区间是: ({[0,0_j], [1 : 1_j], ..[m : m_j]}), 也可以计算 子数组和 小于等于 (Middle) 的子数组的总个数 (Middle_{count}).

注意(): 与「378. 有序矩阵中第K小的元素」十分类似的一个问题就是, 假设我们得到一个 (Middle) ,满足小于等于 (Middle) 的子数组和的个数是 (k) 个, 但是(Middle) 自身并不是子数组和, 举例来说, 还是对于 (nums = [1,2,3,4]). 它的(SUM_{sub} = [1, 2, 3, 3, 4, 5, 6, 7, 9, 10]) , 那么 (Middle) 等于 (8)的时候和 (Middle) 等于 (7) 的时候, (Middle_{count}) 都是 (8), 所以, 我们需要限定二分法找 (Middle) 的时候, 当(Middle_{count}) 大于等于 (k), 那么我们就要缩小 (Middle), 使得 (right = Middle), 而不是 (Middle -1), 具体的思想可以参考「378. 有序矩阵中第K小的元素」.

计算子数组和的前缀和

子数组和的前缀和也就是根据 (SUM_k), 计算前 (k) 大的(SUM_i) 的和, 刚才, 我们已经计算得到了(Middle), 但是可能有多个 (Middle) 大小的子数组和, 所以我们需要将 (sum_{i= 0}^{k}SUM_i) 分成两部分,

一部分是严格小于 (Middle) 的(SUM_i), 这部分的和, 我们刚才分析的时候已经计算过了, 需要注意的时候, 在计算的过程中, 还需要统计 (SUM_i) 的次数, 假设是 (count) 次, 这些子数组和 的和用 (SUM_1) 表示.

另一部分是等于 (Middle) 的(SUM_i), 也就是说 (SUM_i = SUM_k). 这部分的子数组的个数就是 (k-count). 这一部分的结果就是 ((k-count)* Middle).

所以总共的子数组和的和就是: (SUM_1 + (k-count)* Middle). 这一部分就是子数组和的前缀和. 也就是 (sum_{i= 0}^{k}SUM_i) .

那么最后的答案就是 (sum_{i= left}^{right}SUM_i = sum_{i=0}^{right}SUM_i - sum_{i=0}^{left-1}SUM_i)

最后， 贴一下实现的代码：

class Solution {
public:
    static constexpr int MODULO = 1000000007;
    int rangeSum(vector<int>& nums, int n, int left, int right) {
        // 计算我们使用到的前缀和
        vector<int> prefixSums = vector<int>(n);
        prefixSums[0] = nums[0];
        for (int i = 1; i < n; i++) {
            prefixSums[i] = prefixSums[i - 1] + nums[i];
        }
        vector<int> prefixPrefixSums = vector<int>(n);
        prefixPrefixSums[0] = prefixSums[0];
        for (int i = 1; i < n; i++) {
            prefixPrefixSums[i] = prefixPrefixSums[i - 1] + prefixSums[i];
        }
        // 分别计算子数组和排序后的 的前缀和, 第right个 和 第left-1个
        return (getSum(prefixSums, prefixPrefixSums, n, right) - getSum(prefixSums, prefixPrefixSums, n, left - 1)) % MODULO;
    }

    int getSum(vector<int>& prefixSums, vector<int>& prefixPrefixSums, int n, int k) {
        // 如果 k==0 表示left == 1, 相当于left小于最小的子数组和, 所以为 0
        if(k == 0) {
            return 0;
        }
        // 计算第 k 大的子数组和
        int num = getKth(prefixSums, n, k);
        int sum = 0;
        int count = 0;
        // 计算小于 num, 在上述文章中是 Middle 的所有子数组的和
        for (int i = 0, j = 0; i < n; i++) {
            int temp = 0;
            j = i;
            if (i > 0) {
                temp = prefixSums[i - 1];
            }
            while (j < n && prefixSums[j]-temp < num) {
                j++;
            }
            // 只有 j>i 表示子数组的大小至少为 1的时候, 我们加入新的子数组的和
            if(j > i) {
                j--;
                if (i > 0) {
                    sum += prefixPrefixSums[j] - prefixPrefixSums[i - 1] - prefixSums[i - 1] * (1 + j - i);
                }
                else {
                    sum += prefixPrefixSums[j];
                }
                sum %= MODULO;
                count += (j - i + 1);
            }
        }
        // 有 count 个子数组和严格小于 Middle, 剩下的(k-count) 等于 Middle
        sum += num * (k - count);
        return sum;
    }

    int getKth(vector<int>& prefixSums, int n, int k) {
        int low = 0, high = prefixSums[n - 1];
        // 得到第 k 大的子数组和, 注意这里的二分法, 最后返回 low, 
        while (low < high) {
            int mid = (high - low) / 2 + low;
            int count = getCount(prefixSums, n, mid);
            if (count < k) {
                low = mid + 1;
            }
            else {
                high = mid;
            }
        }
        return low;
    }

    int getCount(vector<int>& prefixSums, int n, int x) {
        // 对于给定的 x, 计算他是第几大的子数组和
        int count = 0;
        for (int i = 0, j = 0; i < n; i++) {
            int temp = 0;
            if (i > 0) {
                temp = prefixSums[i - 1];
            }
            while (j < n && prefixSums[j]-temp <= x) {
                j++;
            }
            count += (j - i);
        }
        return count;
    }
};