这篇博客主要是对 327.区间和的个数 - LeetCode官方题解 - 归并排序解法进行更加详细的说明。
参考:LeetCode官方题解
链接:https://leetcode-cn.com/problems/count-of-range-sum/solution/qu-jian-he-de-ge-shu-by-leetcode-solution/
给定一个整数数组
n
u
m
s
nums
nums,返回区间和在
[
l
o
w
e
r
,
u
p
p
e
r
]
[lower, upper]
[lower,upper] 之间的个数,包含
l
o
w
e
r
lower
lower 和
u
p
p
e
r
upper
upper。区间和
S
(
i
,
j
)
S(i, j)
S(i,j) 表示在
n
u
m
s
nums
nums中,位置从
i
i
i 到
j
j
j 的元素之和,包含
i
i
i 和
j
(
i
≤
j
)
j (i ≤ j)
j(i≤j)。
归并排序求解
设数组
n
u
m
s
nums
nums 前缀和数组为
p
r
e
S
u
m
preSum
preSum,
p
r
e
S
u
m
preSum
preSum 定义为:
{
p
r
e
S
u
m
[
0
]
=
0
p
r
e
S
u
m
[
i
+
1
]
=
∑
k
=
0
i
n
u
m
s
[
i
]
(1)
\begin{cases} & preSum[0] = 0 \\ & preSum[i+1]=\sum_{k=0}^{i} nums[i]\\ \end{cases} \tag{1}
{preSum[0]=0preSum[i+1]=∑k=0inums[i](1) 则求区间和个数的问题等价于求所有满足下式的下标对
(
i
,
j
)
(i,j)
(i,j) 的个数:
{
i
<
j
p
r
e
S
u
m
[
j
]
−
p
r
e
S
u
m
[
i
]
∈
[
l
o
w
e
r
,
u
p
p
e
r
]
(2)
\begin{cases} & i<j \\ & preSum[j] - preSum[i] \in [lower, upper] \\ \end{cases} \tag{2}
{i<jpreSum[j]−preSum[i]∈[lower,upper](2) 需要注意在定义
(
1
)
(1)
(1) 中如果不加上
p
r
e
S
u
m
[
0
]
=
0
preSum[0] = 0
preSum[0]=0 的话,数组
n
u
m
s
nums
nums 第一个元素
n
u
m
s
[
0
]
nums[0]
nums[0] 就不能用公式
p
r
e
S
u
m
[
j
]
−
p
r
e
S
u
m
[
i
]
preSum[j] - preSum[i]
preSum[j]−preSum[i] 表示,从而
n
u
m
s
[
0
]
nums[0]
nums[0] 单独构成的区间就检测不到。
如果直接用循环求解的话时间复杂度还是 O ( n 2 ) O(n^2) O(n2),可能比暴力解法还要麻烦,我们转化为前缀和数组的形式主要是为了应用归并排序的思想。
先考虑如下的问题:给定两个升序排列的数组
n
1
,
n
2
n1, n2
n1,n2,找出所有满足下式的下标对
(
i
,
j
)
(i, j)
(i,j):
n
2
[
j
]
−
n
1
[
i
]
∈
[
l
o
w
e
r
,
u
p
p
e
r
]
(3)
n2[j]-n1[i] \in [lower, upper] \tag{3}
n2[j]−n1[i]∈[lower,upper](3) 首先我们在数组
n
2
n2
n2 中设置两个指针
l
,
r
l,r
l,r,它们都指向
n
2
n2
n2 的起始位置。然后执行如下流程:
首先考察 n 1 n1 n1 第一个元素 n 1 [ 0 ] n1[0] n1[0],不断向右移动指针 l l l,直到 n 2 [ l ] ≥ n 1 [ 0 ] + l o w e r n2[l] \geq n1[0] + lower n2[l]≥n1[0]+lower,然后向右移动指针 r r r,直到 n 2 [ r ] > n 1 [ 0 ] + u p p e r n2[r] > n1[0] + upper n2[r]>n1[0]+upper,这样对于 i = 0 , j ∈ [ l , r ) i=0,j \in [l,r) i=0,j∈[l,r) 中的坐标都满足公式 ( 3 ) (3) (3),
接着我们考察 n 1 n1 n1 第二个元素 n 1 [ 1 ] n1[1] n1[1],因为两个数组都是升序的,所以 n 1 [ 1 ] > n 1 [ 0 ] n1[1] > n1[0] n1[1]>n1[0],并且因为 n 2 [ l ] n2[l] n2[l] 是数组 n 2 n2 n2 中第一个 ≥ n 1 [ 0 ] + l o w e r \geq n1[0] + lower ≥n1[0]+lower的的数,所以我们只需要从指针 l l l 当前位置开始就可以找到第一个 ≥ n 1 [ 1 ] + l o w e r \geq n1[1] + lower ≥n1[1]+lower的的数,不断向右移动指针 l l l,直到 n 2 [ l ] ≥ n 1 [ 1 ] + l o w e r n2[l] \geq n1[1] + lower n2[l]≥n1[1]+lower,同理我们也只需从当前位置向右移动指针 r r r,直到 n 2 [ r ] > n 1 [ 1 ] + u p p e r n2[r] > n1[1] + upper n2[r]>n1[1]+upper,最后得到满足公式 ( 3 ) (3) (3) 的坐标集 i = 1 , j ∈ [ l , r ) i=1,j \in [l,r) i=1,j∈[l,r)。
不断重复上述过程,过程中指针 l , r l,r l,r 只能向右移动,对于 n 1 n1 n1 中每一个下标 i i i,都可以得出 n 2 n2 n2 下标 j j j 满足条件区间 [ l , r ) [l,r) [l,r) 的大小,最终可以得到所有满足条件的下标对 ( i , j ) (i,j) (i,j) 的数量,这样我们解决了两个升序数组满足条件的坐标对数量的问题。
在归并排序中,先将一个待排序的数组分解至若干个只有一个元素的数组,然后不断两两合并,要求合并后的数组是有序的,即每次将两个有序的数组合成一个有序的数组,这个过程也可以用两个数组中不断向右移动的两个指针实现。(p.s.我描述的可能不太准确,不懂得可以看归并排序 - 知乎)。
最后我们就可以解决求数组区间和个数这个问题了,整个解法过程为:在对前缀和数组进行归并排序的过程中,每次合并两个数组的同时计算合并后数组中满足条件下标对的数目。
至于合并后数组满足条件下标对的数目怎么计算,官方题解中给出:合并后数组满足条件的下标对数目 = = = 合并前两个数组满足条件的下标对数目之和 + + + 左端点在左侧数组同时右端点在右侧数组的下标对数量。
因为这是一个递归过程,当数组只有一个元素时(递归最底层)它满足条件的下标对数目为 0 0 0,然后根据上面的公式不断递归下去就 O K OK OK 了,那么为什么是 + + + 左端点在左侧数组同时右端点在右侧数组的下标对数量?
将合并前的两个子数组的时候根据下标顺序分成左侧数组和右侧数组,由于合并后的数组满足要求的下标对中下标
i
,
j
i,j
i,j 都只在左侧数组和都只在右侧数组的下标对数目我们已经得到了,所以只需要算出下标
i
,
j
i,j
i,j 分别在不同数组的下标对数目就行了,又因为下标对
(
i
,
j
)
(i,j)
(i,j) 需要满足
i
<
j
i<j
i<j,所以下标
i
i
i 只能在左侧数组中取,下标
j
j
j 只能在右侧数组中取,这样就是求两个升序数组
n
[
1
]
,
n
[
2
]
n[1],n[2]
n[1],n[2] 满足条件
(
3
)
(3)
(3) 的下标对数目了,将左侧数组作为
n
[
1
]
n[1]
n[1],右侧数组作为
n
[
2
]
n[2]
n[2],按照前面说的方法就可以求解,这样整个问题就解决了。
Python实现:
class Solution327:
def countRangeSum(self, nums, lower, upper):
sum_nums = [0]
for i in range(1, len(nums) + 1):
sum_nums.append(sum_nums[i-1] + nums[i-1])
return self.countRangeSumRecursive(sum_nums, lower, upper, 0, len(sum_nums) - 1)
def countRangeSumRecursive(self, nums, lower, upper, left, right):
if left == right:
return 0
mid = (left + right) // 2
count_L = self.countRangeSumRecursive(nums, lower, upper, left, mid)
count_R = self.countRangeSumRecursive(nums, lower, upper, mid+1, right)
count = count_L + count_R
# 计算下标对数目
l = r = mid + 1
for i in range(left, mid+1, 1):
while l <= right and nums[l] - nums[i] < lower:
l += 1
while r <= right and nums[r] - nums[i] <= upper:
r += 1
count += r - l
# 归并排序部分
sorted_nums = [0] * (right - left + 1)
l, r, i = left, mid+1, 0
while l <= mid and r <= right:
if nums[l] <= nums[r]:
sorted_nums[i] = nums[l]
l += 1
else:
sorted_nums[i] = nums[r]
r += 1
i += 1
sorted_nums[i:] = nums[l:mid+1] if r == right+1 else nums[r:right+1]
nums[left:right+1] = sorted_nums
return count