子数组和排序后的区间和

标签: 数组 双指针 二分查找 排序

难度: Medium

给你一个数组 nums ，它包含 n 个正整数。你需要计算所有非空连续子数组的和，并将它们按升序排序，得到一个新的包含 n * (n + 1) / 2 个数字的数组。

请你返回在新数组中下标为 left 到 right （下标从 1 开始）的所有数字和（包括左右端点）。由于答案可能很大，请你将它对 10^9 + 7 取模后返回。

示例 1：

输入：nums = [1,2,3,4], n = 4, left = 1, right = 5
输出：13 
解释：所有的子数组和为 1, 3, 6, 10, 2, 5, 9, 3, 7, 4 。将它们升序排序后，我们得到新的数组 [1, 2, 3, 3, 4, 5, 6, 7, 9, 10] 。下标从 le = 1 到 ri = 5 的和为 1 + 2 + 3 + 3 + 4 = 13 。

示例 2：

输入：nums = [1,2,3,4], n = 4, left = 3, right = 4
输出：6
解释：给定数组与示例 1 一样，所以新数组为 [1, 2, 3, 3, 4, 5, 6, 7, 9, 10] 。下标从 le = 3 到 ri = 4 的和为 3 + 3 = 6 。

示例 3：

输入：nums = [1,2,3,4], n = 4, left = 1, right = 10
输出：50

提示：

  • 1 <= nums.length <= 10^3
  • nums.length == n
  • 1 <= nums[i] <= 100
  • 1 <= left <= right <= n * (n + 1) / 2

Submission

运行时间: 62 ms

内存: 16.3 MB

# binary search + preprefix sum
from typing import List
class Solution:
    def rangeSum(self, nums: List[int], n: int, left: int, right: int) -> int:
        def count(x):
            # search how many numbers smaller than x in 2d sorted array
            cnt = 0
            j = 1
            for i in range(n):
                # prefix should include the j == n
                # lower bound but also count the same numbers
                while j <= n and pre[j] - pre[i] <= x:
                    j += 1
                cnt += j-i-1
            return cnt

        def pre_sum(k):
            # binary search the kth small number
            low = 0
            high = pre[-1]
            while low < high:       # lower bound
                mid = (low+high)//2
                if count(mid) >= k:
                    high = mid
                else:
                    low = mid+1
            k_small = low

            # get sum
            res = 0
            cnt = 0
            j = 1
            for i in range(n):
                while j <= n and pre[j]-pre[i] < k_small:
                    j += 1
                j -= 1
                res += prepre[j] - prepre[i] - pre[i]*(j-i)
                cnt += j-i
            # add numbers that equal to kth small number
            return res + k_small*(k-cnt)

        MOD = int(1e9+7)
        pre = [0]
        for num in nums:
            pre.append(pre[-1] + num)
        prepre = [0]
        for num in pre[1:]:
            prepre.append(prepre[-1] + num)
        return (pre_sum(right) - pre_sum(left-1)) % MOD

print(Solution().rangeSum([7,5,8,5,6,4,3,3],8,2,6))

Explain

这个题解使用了二分查找和前缀和预处理的方法来解决问题。首先，对原数组进行前缀和预处理，得到pre数组和prepre数组，分别表示原数组的前缀和以及前缀和数组的前缀和。然后，通过二分查找来确定第k小的子数组和，同时统计比第k小的数小的子数组和以及等于第k小的数的子数组和，最后将它们相加得到结果。在二分查找的过程中，使用了count函数来计算二维有序数组中小于等于某个数的元素个数。

时间复杂度: O(n * log(range))

空间复杂度: O(n)

from typing import List
class Solution:
    def rangeSum(self, nums: List[int], n: int, left: int, right: int) -> int:
        def count(x):
            # 在二维有序数组中搜索小于等于x的数的个数
            cnt = 0
            j = 1
            for i in range(n):
                # 前缀和应该包括 j == n 的情况
                # 下界，同时也要计算相等的数
                while j <= n and pre[j] - pre[i] <= x:
                    j += 1
                cnt += j-i-1
            return cnt

        def pre_sum(k):
            # 二分查找第k小的数
            low = 0
            high = pre[-1]
            while low < high:       # 下界
                mid = (low+high)//2
                if count(mid) >= k:
                    high = mid
                else:
                    low = mid+1
            k_small = low

            # 计算和
            res = 0
            cnt = 0
            j = 1
            for i in range(n):
                while j <= n and pre[j]-pre[i] < k_small:
                    j += 1
                j -= 1
                res += prepre[j] - prepre[i] - pre[i]*(j-i)
                cnt += j-i
            # 加上等于第k小的数的子数组和
            return res + k_small*(k-cnt)

        MOD = int(1e9+7)
        pre = [0]
        for num in nums:
            pre.append(pre[-1] + num)
        prepre = [0]
        for num in pre[1:]:
            prepre.append(prepre[-1] + num)
        return (pre_sum(right) - pre_sum(left-1)) % MOD

Explore

在二分查找中，`low`的初始值设为0，因为子数组和不会小于0（考虑到所有元素都非负的情况）。`high`的初始值设为`pre[-1]`，即所有元素的总和，这是因为子数组和的最大可能值就是整个数组的和。

在`count`函数中，`j`的初始值为1是因为我们需要计算的是从第i个元素开始的子数组和，即`pre[j] - pre[i]`形式。如果`j`从0开始，那么`pre[i] - pre[0]`将计算从数组开始到第i个元素的和，这不符合我们计算从i到j的子数组和的需求。

`pre`数组存储了从数组开始到当前位置的累计和。因此，`pre[1:]`中的每个元素`num`实际上代表从数组开始到当前位置的总和。`prepre`数组是`pre`数组的前缀和，这意味着`prepre`中的每个元素是`pre`数组中所有前面元素的总和。所以，`prepre.append(prepre[-1] + num)`实际上是在累计`pre`数组的值，从而用来快速计算任意两个索引i和j之间的前缀和之差的总和。

`count`函数通过滑动窗口的方式来确定满足条件的子数组个数。为了确保边界的准确性，函数内部使用循环和条件判断来精确地控制索引`j`的移动。当`j`不能再向前移动时（即`pre[j] - pre[i]`大于x），循环停止，这保证了不会错过任何可能的子数组。同时，每次循环都会根据不同的i值重新调整j的位置，以确保覆盖所有可能的子数组结构。