题目描述
给定一个长度为 n 的数组 a（元素可正可负可为0），求有多少个非空连续子段 [l, r]\ 满足子段和 $$a_l + a_{l+1} + \dots + a_r < t$$
输入格式
第一行两个整数 n 和 \(t\)（\(1 \le n \le 200\,000\)，\(|t| \le 2 \times 10^{14}\)）
第二行 (n) 个整数 \(a_i\)(\(|a_i| \le 10^9\)）
输出格式
输出满足条件的子段个数
解题思路
前缀和转化
设前缀和数组 \(pre[i] = \sum_{k=1}^{i} a_k\)，特别地 \(pre[0] = 0\)
区间 \([l, r]\) 的和可表示为 \(pre[r] - pre[l-1]\)
条件 $$pre[r] - pre[l-1] < t$$ 等价于
\[pre[l-1] > pre[r] - t.
\]
固定右端点 \(r\)，我们需要统计有多少个左端点 \(l\)\((1 \le l \le r\)）满足上式，即有多少个下标 \(j = l-1\)（\(0 \le j \le r-1\)）使得 \(pre[j] > pre[r] - t\)
因此，我们可以从左到右扫描数组，维护一个数据结构，其中已经插入了 \(pre[0], pre[1], \dots, pre[r-1]\)。对于当前 (r)，查询数据结构中值大于 (pre[r] - t) 的元素个数，累加到答案，然后将 (pre[r]) 插入数据结构
数据结构选择
需要支持动态插入和查询“大于某值的元素个数”
- 平衡树（如
std::multiset）配合upper_bound可求大于某值的个数，但需要手写距离（可用distance，但复杂度可能退化） - 更优雅的方式：使用基于红黑树的
__gnu_pbds::tree，它提供了order_of_key方法，返回严格小于给定键的元素个数
我们只需插入所有前缀和（由于可能有重复值，用pair<ll, ll>存储，第二维为下标保证唯一性），然后对于给定的need = pre[r] - t，\[\text{大于 need 的个数} = \text{总个数} - \text{order_of_key}(\{need, INF\}). \]其中{need, INF}是一个大于所有值为need的元素的键，这样order_of_key返回的就是小于等于need的元素个数
算法步骤
- 读入 (n, t) 和数组 (a)
- 计算前缀和 (pre[i])
- 初始化一个
ordered_set，类型为tree<pair<ll, ll>, ...>，插入{pre[0], 0} - 遍历 (i = 1 \dots n)：
- 计算
need = pre[i] - t - 查询
cnt = st.order_of_key({need, LLONG_MAX})，表示小于等于need的元素个数 - 当前答案加上
st.size() - cnt，即为大于need的元素个数 - 插入
{pre[i], i}
- 计算
- 输出答案
复杂度分析
- 时间复杂度：每个前缀和插入一次、查询一次，每次操作 \(O(\log n)\)，总复杂度 \(O(n \log n)\)
- 空间复杂度：\(O(n)\)
代码实现（C++）
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using ll = long long;
using namespace __gnu_pbds;
using ordered_set = tree<pair<ll, ll>, null_type, less<pair<ll, ll>>, rb_tree_tag, tree_order_statistics_node_update>;
const ll N = 2e5 + 10, mod = 998244353, inf = LLONG_MAX;
ll n, w, m, t, a[N];
void solve() {cin >> n >> t;ll ans = 0;vector<ll> pre(n + 2, 0);for (int i = 1; i <= n; i++) {cin >> a[i];pre[i] = pre[i - 1] + a[i];}ordered_set st;st.insert({pre[0],0});//放入下表为0的值for (int i = 1; i <= n; i++) {ll need=pre[i]-t;ll nn=st.order_of_key({need,inf});ans += st.size()-nn;//大于need的个数st.insert({pre[i],i});}cout << ans << '\n';
}
int main() {ios::sync_with_stdio(0);cin.tie(0);cout.tie(0);int _ = 1;// cin >> _;while (_--) {solve();}return 0;
}