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;
    }
};