【Leetcode】1508. Range Sum of Sorted Subarray Sums

题目地址：

https://leetcode.com/problems/range-sum-of-sorted-subarray-sums/description/

给定一个长$n$的正整数数组$a$，如果将其所有子数组和求出来，则一共有$\frac{n(n+1)}{2}$个。将这些和从小到大排序，问排序之后位于第$l$到$r$个的那些和的总和是多少。

设子数组的和组成的数组是$b$且已经排好序。先预处理一下$a$的前缀和数组$s$，即$s[i]$是$a$的前$i$个数的和，从$1$开始计数。我们考虑如何求$b$中排名第$l$的数，排名第$l$的数$x$其实是小于$x$的数的个数小于$l$的最大的数，而这个可以用二分答案来做。求完了排名第$l$的数$x$以后，由于$\sum b[l:r]=\sum b[1:r]-\sum b[1:l-1]$，所以我们只需要会求$f(k)=\sum b[1:k]$即可。我们可以在$b$里先求小于$x$的数的个数$c$和小于$x$的数的总和$S$，那么$f(l)=S+(l-c)x$，而在$b$里先求小于$x$的数的个数$c$和小于$x$的数的总和$S$这一步，可以在前缀和$s$数组上用双指针来求。固定$i$，我们找到最左边的$j$使得$\sum s[j:i]<x$，由于$a$只含正整数，所以$s[.\ge j:i]<x$，即以$i$为右端点且和小于$x$的子段个数为$i-j+1$个，设前缀和数组$s$自己的前缀和数组为$ss$（$ss$长度也是$n+1$，且$ss[i]=\sum s[1:i]$），则小于$x$的所有子段和的和就是$$\begin{gathered} (s[i]-s[j])+...+(s[i]-s[i-1]) \\ =(i-j)s[i]-(ss[i-1]-ss[j-1]) \end{gathered}$$代码如下：

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 \sum a)$，空间$O(1)$。