题目地址:
https://leetcode.com/problems/range-sum-of-sorted-subarray-sums/description/
给定一个长 n n n的正整数数组 a a a,如果将其所有子数组和求出来,则一共有 n ( n + 1 ) 2 \frac{n(n+1)}{2} 2n(n+1)个。将这些和从小到大排序,问排序之后位于第 l l l到 r r r个的那些和的总和是多少。
设子数组的和组成的数组是 b b b且已经排好序。先预处理一下 a a a的前缀和数组 s s s,即 s [ i ] s[i] s[i]是 a a a的前 i i i个数的和,从 1 1 1开始计数。我们考虑如何求 b b b中排名第 l l l的数,排名第 l l l的数 x x x其实是小于 x x x的数的个数小于 l l l的最大的数,而这个可以用二分答案来做。求完了排名第 l l l的数 x x x以后,由于 ∑ b [ l : r ] = ∑ b [ 1 : r ] − ∑ b [ 1 : l − 1 ] \sum b[l:r]=\sum b[1:r]-\sum b[1:l-1] ∑b[l:r]=∑b[1:r]−∑b[1:l−1],所以我们只需要会求 f ( k ) = ∑ b [ 1 : k ] f(k)=\sum b[1:k] f(k)=∑b[1:k]即可。我们可以在 b b b里先求小于 x x x的数的个数 c c c和小于 x x x的数的总和 S S S,那么 f ( l ) = S + ( l − c ) x f(l)=S+(l-c)x f(l)=S+(l−c)x,而在 b b b里先求小于 x x x的数的个数 c c c和小于 x x x的数的总和 S S S这一步,可以在前缀和 s s s数组上用双指针来求。固定 i i i,我们找到最左边的 j j j使得 ∑ s [ j : i ] < x \sum s[j:i]<x ∑s[j:i]<x,由于 a a a只含正整数,所以 s [ . ≥ j : i ] < x s[.\ge j:i]<x s[.≥j:i]<x,即以 i i i为右端点且和小于 x x x的子段个数为 i − j + 1 i-j+1 i−j+1个,设前缀和数组 s s s自己的前缀和数组为 s s ss ss( s s ss ss长度也是 n + 1 n+1 n+1,且 s s [ i ] = ∑ s [ 1 : i ] ss[i]=\sum s[1:i] ss[i]=∑s[1:i]),则小于 x x x的所有子段和的和就是 ( s [ i ] − s [ j ] ) + . . . + ( s [ i ] − s [ i − 1 ] ) = ( i − j ) s [ i ] − ( s s [ i − 1 ] − s s [ j − 1 ] ) (s[i]-s[j])+...+(s[i]-s[i-1])\\=(i-j)s[i]-(ss[i-1]-ss[j-1]) (s[i]−s[j])+...+(s[i]−s[i−1])=(i−j)s[i]−(ss[i−1]−ss[j−1])代码如下:
class Solution {
public:
using ll = long long;
using PLL = pair<ll, ll>;
int rangeSum(vector<int>& a, int n, int l, int r) {
constexpr ll MOD = 1e9 + 7;
vector<ll> s(n + 1), ss(n + 1);
// 求a的前缀和数组s,和s的前缀和数组ss,其中ss我们不考虑空前缀,直接让ss长度也是n + 1
for (int i = 1; i <= n; i++) {
s[i] = s[i - 1] + a[i - 1];
ss[i] = ss[i - 1] + s[i];
}
// 求和小于x的子段个数和它们的总和
auto get_sum_and_cnt = [&](int x) -> PLL {
ll sum = 0, cnt = 0;
for (int i = 1, j = 0; i <= n; i++) {
while (j < i && s[i] - s[j] >= x) j++;
if (j < i) {
cnt += i - j;
sum += (i - j) * s[i] - (ss[i - 1] - (j ? ss[j - 1] : 0));
}
}
return {sum, cnt};
};
// 求排名为k的子段和
auto get_kth = [&](int k) -> int {
int l = 0, r = s[n];
while (l < r) {
int mid = l + (r - l + 1 >> 1);
auto [sum, cnt] = get_sum_and_cnt(mid);
if (cnt < k)
l = mid;
else
r = mid - 1;
}
return l;
};
auto f = [&](int k) -> ll {
auto x = get_kth(k);
auto [sum, cnt] = get_sum_and_cnt(x);
return sum + (k - cnt) * x;
};
return (f(r) - f(l - 1)) % MOD;
}
};
时间复杂度 O ( n log ∑ a ) O(n\log \sum a) O(nlog∑a),空间 O ( 1 ) O(1) O(1)。