<二分查找+双指针+前缀和>解决子数组和排序后的区间和
题目重现:
给你一个数组 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
来源:力扣(LeetCode)
链接:https://leetcode-cn.com/problems/range-sum-of-sorted-subarray-sums
这是在leetcode上碰到的一道题,但由于设置的测试样例并不是很好,而导致暴力解法也可通过,所以此题只是中等难度。但看过题解的解法思路后觉得有必要做一记录。由浅入深,先通过暴力解法,然后引出优化的方法。
暴力法
这道题给出一个数组nums,如果暴力解题,可以先计算出它的所有非空连续子数组的和,然后进行排序,再计算它下标left到right的和,最后取余数即可。
列举出所有的非空连续子数组,使用左右双指针,假设题目给定nums为1,2,3,4
,那么先让左指针指1,右指针从1开始依次滑动过整个数组后面的数,即可得到以1开头的子数组和为1,3,6,10
,再让左指针右移一位,继续按上述可得2,5,9
......以此类推可得所有子数组,然后对其进行排序。子数组和的数目总共为n*(n+1)/2
个。
//暴力法
class Solution {
public int rangeSum(int[] nums, int n, int left, int right) {
int[] new_arr = new int[n*(n+1)/2+1]; //定义数组存放所有子数组
int index = 1;
for (int i = 0; i < nums.length; i++) {
int pre = 0;
for (int j = i; j < nums.length; j++) {
new_arr[index++] = pre+nums[j]; //dp思想,左指针固定后,右指针滑动后的下一个子数组等于上次加nums[j]之和
pre = pre+nums[j]; //更新pre
}
}
Arrays.sort(new_arr); //对子数组进行排序
long count = 0;
for (int i = left; i <= right; i++) {
count+=new_arr[i];
}
while (count >= 1000000007) { //取余数后返回
count -= 1000000007;
}
return (int)count;
}
}
前置讨论
讨论二分查找+双指针解法前,先看leetcode的另一道题378. 有序矩阵中第K小的元素,这道题的解题思路有助于我们更好的解决上面的题目。
给定一个 n x n 矩阵,其中每行和每列元素均按升序排序,找到矩阵中第 k 小的元素。
请注意,它是排序后的第 k 小元素,而不是第 k 个不同的元素。示例:
matrix = [
[ 1, 5, 9],
[10, 11, 13],
[12, 13, 15]
],
k = 8,返回 13。
来源:力扣(LeetCode)
链接:https://leetcode-cn.com/problems/kth-smallest-element-in-a-sorted-matrix
先观察这个给定的数组matrix发现:整个数组的行从左到右递增,从上到下递增,刚开始我想的是用优先级队列,首先加入一个最小的数(最左上角),然后每次加入队列头的右边的数和下边的数,周而复始的循环k次,队列头就是这个数。但由于优先级队列的维护本身就是非常耗时的,所以整个程序执行下来时间效率很低,运行了42ms。下面给出优化思路:
二分查找的思路,以下图为例:(图片来自leetcode官方题解)
通过观察发现mid = (1+16)/2 = 8,大于mid的都分布在红线下面,而不大于mid的部分分布在红线上面,所以可以使用二分查找。
沿着图中蓝色箭头走一边,就可以计算出上方板块的大小,即不大于mid的数字的数目,这样通过二分将mid逐渐逼近第k小的元素。
算法描述:目的是统计不大于当前mid的数的数目,从第0列最后一行开始,如果此列最下面的数都不大于mid,那么此列所有的数肯定都不大于mid,继续到下一列,将列指针向右移动,如果此时最后一行的数大于mid,则将指示行的指针上移直到遇到一个不大于mid的数就停止,而这个数上面的数肯定都不大于mid。如果行指针已经滑到0还没有不大于mid的数出现,那说明后面已经不可能有不大于mid的数了,因为这个数组向右和向下是递增的。
当访问第j列的时候,如果第i+1行大于mid,而第i行不大于mid,则这列不大于mid的数数目为i+1(考虑第0行)。统计整个数组中不大于mid的数的数目。如果小于k,则说明mid太小,将left右移至mid+1处,否则将right移至mid处。直到左右指针相遇,此时它们所指向的就是第k小的数。
private static int kthSmallest(int[][] matrix, int k) {
int n = matrix.length;
int left = matrix[0][0];
int right = matrix[n-1][n-1];
int mid;
//二分查找,找到第k小的数
while (left < right) {
mid = left + ((right-left) >> 1);
if (check(matrix,mid,k,n)) {
right = mid;
} else {
left = mid + 1;
}
}
return left;
}
//利用双指针检查当前mid是否过大(即是否在数组matrix中比mid大的数超过了k个)
private static boolean check (int[][] matrix, int mid, int k, int n) {
int i = n-1; //指示行坐标
int j = 0; //指示列坐标
int num = 0;
while (i >= 0 && j < n) {
if (matrix[i][j] <= mid) {
j++;
num += (i+1);
} else {
i--;
}
}
return num>=k;
}
Q1:考虑check
函数,为什么要matrix[i][j] <= mid
而不是matrix[i][j] < mid
?
A:因为数组matrix中可能会出现重复的数字,加入第k小的数字也重复了,形如1,2,3,3,3,3,3,5,6
第5小的数字,那第3和第4个数字等于3,当却确实在第5小的数字之前。
Q2:考虑check
函数,为什么要return num >= k
而不是return num > k
?
A:因为left和right以及mid是从最小到最大的数之间的任意一个数,所以并不能保证它们就一定是数组中存在的数,如果某个mid能保证小于等于它的数恰好为k个,则第k小的数就是它之前最近的一个存在于数组中的数。所以当此时不大于mid的数大于或 等于k个时,就可以保证要求的第k小的数
一定在mid或mid之前,故而将right移动到mid处。
Q3:考虑kthSmallest
函数,为什么check
函数返回为真就左移,假就右移?
A:设置第k小的数为res,当mid在res左边时,此时数组中不大于mid的数会因为少了res而小于k,因为res是第k个,所以left会右移,以求使mid右移。当mid在res右边时或者就刚好等于res时,此时数组中大于等于mid的数会等于或超过k个,由于res为第k个,而mid又在其右或等于,所以此时数组中不大于mid的数至少为k个,所有使right左移,使得mid左移。
Q4:考虑kthSmallest
函数,为什么能保证最后left指向的就是数组中的元素呢?
A:根据check
函数返回的情况不断将左右指针逼近res,mid总是在res左右横跳,带动left和right逼近res,而mid终会有一次等于res,此时不大于mid的数大于或等于k个,右指针左移到res上,这时的mid总是小于res,而导致不大于mid的数目小于k,左指针右移,左右指针相邻时,下一次左指针移动,必定移动到res上,left == right
,跳出循环。
Q5:考虑kthSmallest
函数,为什么right = mid
,而left = mid+1
呢?
A:这是由于除法向下取整而导致的二分查找的特性,假设此时left=2,right=3
,则mid=2
,如果left = mid
,则会一直原地打转。如果right = mid-1
,则可能此时的mid == res
(mid==res时必定是右指针左移,参考Q3),右指针就会移动到res之前,从而错过正解。
优化解法
继续回到这个题,看完前面前置的讨论后相信对解答这个题会有很大帮助。如题目给的示例1:nums = {1,2,3,4}
,这样我们可以构造出它的非空连续子数组的和矩阵如下:
第1行是以1开头的子数组的和,分别对应1;1,2;1,2,3;1,2,3,4
,第2行是以2行开始的子数组的和,以此类推,观察此数组发现,这个数组从左到右以此递增,从上到下以此递增,看到这应该就明白了上面那个前置讨论的意义了。
先确定我们的大思路:题目要求构造一个非空连续子数组的和,在这个新数组中从left到right的元素之和,那我们可以参考前置讨论里的方法先得到前left-1
大的数字,然后计算前left-1
个数字之和记为f(left-1)
,再同理计算前right个数字之和记为f(right)
,最后答案就是f(right) - f(left-1)
。
<h5 id="1">flag</h5>
计算第k小的数字时候构造以1开始的前缀和数组sums
,数组大小为n+1,我们实际有意义的从1开始,数组的第0个初始化为0,这样就不需要构建整个二位数组了,而计算第2行的时候,发现第2行的每列数字对应上一行相应列的数字只是少了sums[1]
,第三行相比第一行来说就是少了sums[2]
,所以只需用第一行的数字依次减去sums[i]
就是第i行的各数,比如第二行的5就等于sums[3] - sums[1] = 6 - 1
。
/**
* 获取小于mid的数的个数
* @param sums 原数组的前缀和
* @param n 原数组的大小
* @param mid 二分法中的当前mid
* @return 返回严格小于mid数的个数
*/
private int getCnt (long[] sums, int n, int mid) {
int res = 0; //返回的个数
for (int i = 0, p = 1; i < n; i++) {
while (p <= n && sums[p] - sums[i] <= mid) {
p++;
}
//因为每次符合都对p++,所以当最后一次符合条件后也对p进行了加1操作,而加1后p已经指向了最后一个符合条件的下一个数,所以还要给p-1
res += p-1-i;
}
return res;
}
/**
* 利用二分查找获取第k小的数
* @param sums 原数组的前缀和
* @param n 原数组的大小
* @param k 第k小
* @return 返回第k小的数
*/
private int getKth (long[] sums, int n, int k) {
int left = 0, right = Integer.MAX_VALUE; //二分查找指示左右的两个指针
while (left < right) {
int mid = left + ((right-left) >> 1);
if (getCnt(sums, n, mid) >= k) {
right = mid;
} else {
left = mid + 1;
}
}
return left;
}
我们设计一个getSum(k)
这个函数,就是上述的f函数,用来计算前k小的数字之和,计算时我们使用前缀和数组,并构造一个前缀和的前缀和数组,如下示例:
此时我们已经得到了第k小的数字,要计算前k小的所有数字之和,考虑到第k小的数字会有重复大小的数字,所以分开计算,明确一点:我们已经得到第k小个数字,假设为6,前k小的数字为1,2,3,6,6,6
,可能后面还有几个6,不过由于k个数的限制,并不纳入计算,所以我们先计算严格小于6的数字之和以及这些数字的个数记为cnt,然后加上(k-cnt) * 6
。
我们构造出了前缀和数组sums和前缀和的前缀和数组ssums
这样以来如果我们要计算第1行的sums[2]+sums[3]
的和,由于ssums[3] = sums[1] + sums[2] + sums[3]
,而ssums[1] = sums[1]
,所以sums[2] + sums[3] = ssums[3] - ssums[1]
。
但是,我们如果要计算第2行的2+5要如何计算呢,通过前面的ssums[3] - ssums[1] - 2*1
,其实整个第2行都会比第1行少1,而第i行会比第1行少nums[i]
。
因此对于连续非空子数组的和构成的数组我们要求所有严格小于第k小的数(记为kth)的和,遍历每一行,每行都是从小到大递增,当找到此行比kth小的最后一个数后,只需要根据前缀和的前缀和数组就可在O(1)的时间里算出来,假设第i行的第p列是此行最后一个小于kth的数,则此行小于kth的数字和为ssums[p] - ssum[i] - (p-i)*nums[i]
。
class Solution {
final int MODULO = 1000000007;
//二分+双指针
/**
* 获取小于mid的数的个数
* @param sums 原数组的前缀和
* @param n 原数组的大小
* @param mid 二分法中的当前mid
* @return 返回严格小于mid数的个数
*/
private int getCnt (long[] sums, int n, int mid) {
int res = 0; //返回的个数
for (int i = 0, p = 1; i < n; i++) {
while (p <= n && sums[p] - sums[i] <= mid) {
p++;
}
//因为每次符合都对p++,所以当最后一次符合条件后也对p进行了加1操作,而加1后p已经指向了最后一个符合条件的下一个数,所以还要给p-1
res += p-1-i;
}
return res;
}
/**
* 利用二分查找获取第k小的数
* @param sums 原数组的前缀和
* @param n 原数组的大小
* @param k 第k小
* @return 返回第k小的数
*/
private int getKth (long[] sums, int n, int k) {
int left = 0, right = Integer.MAX_VALUE; //二分查找指示左右的两个指针
while (left < right) {
int mid = left + ((right-left) >> 1);
if (getCnt(sums, n, mid) >= k) {
right = mid;
} else {
left = mid + 1;
}
}
return left;
}
/**
* 获取前k小的数的和
* @param sums 原数组的前缀和
* @param ssums 原数组前缀和的前缀和
* @param n 原数组大小
* @param k k
* @return 返回前k小的数字之和
*/
private long getSum (long[] sums, long[] ssums, int n, int k) {
long res = 0, cnt = 0;
long kth = getKth(sums, n, k); //第k小的数字
//分两部分计算,考虑到有的数字会重复,所以先计算严格小于kth的数字的和与个数cnt,在加上剩余k-cnt个第k小的数字
for (int i = 0, p = 1; i < n; i++) {
while (p<=n && sums[p]-sums[i] < kth) {
p++;
}
res = (res + ssums[p-1] - ssums[i] - (long)(p-1-i)*sums[i]);
cnt += p-1-i;
}
return (res + (k-cnt)*kth);
}
/**
* 计算
* @param nums
* @param n
* @param left
* @param right
* @return
*/
public int rangeSum (int[] nums, int n, int left, int right) {
long[] sums = new long[n+1];
long[] ssums = new long[n+1];
for (int i = 1; i <= n; i++) {
sums[i] = sums[i-1]+nums[i-1];
ssums[i] = ssums[i-1]+sums[i];
}
long r = getSum(sums, ssums, n, right);
long l = getSum(sums, ssums, n, left-1);
return (int) ((r-l)%MODULO);
}
}