1. 问题描述:
给定一个长度为 n 的整数数列 a1,a2,…,an 和一个整数 t,请你判断共有多少个数对 (l,r) 同时满足:
1 ≤ l ≤ r ≤ n
al + al+1+ ... + ar−1 + ar < t
输入格式
第一行包含两个整数 n 和 t。第二行包含 n 个整数 a1,a2,…,an。
输出格式
一个整数,表示满足条件的数对的数量。
数据范围
前三个测试点满足 1 ≤ n ≤ 5;
所有测试点满足 1 ≤ n ≤ 2 × 10 ^ 5,|t| ≤ 2 × 10 ^ 14,|ai| ≤ 10 ^ 9。
输入样例1:
5 4
5 -1 3 4 -1
输出样例1:
5
输入样例2:
3 0
-1 2 -3
输出样例2:
4
输入样例3:
4 -1
-2 1 -2 3
输出样例3:
3
来源:https://www.acwing.com/problem/content/description/4319/
2. 思路分析:
分析题目可以知道我们需要求解所有满足条件的区间个数,所以首先需要想的是如何枚举才可以不重不漏地将所有答案枚举出来,这里既可以枚举左端点也可以枚举右端点,这里枚举的是右端点这样比较自然一点,也即枚举右端点i,以 i 为右端点求解有多少个区间满足[j,i] 使得si - sj-1 < t,相当于是求解以 i 为右端点 的前面有多少个 j 满足 sj > si - t,这样就可以不重不漏地将所有答案枚举出来,可以发现我们需要动态维护到当前位置的有序序列,并且需要求解出 sj 的排名,对于动态维护有序序列可以使用平衡树来解决(c++可以使用 set),但是使用 set 就不能够求解排名了,但是新版的 c++ 提供了 tree 可以实现这个功能;为了写代码方便处理除了使用平衡树来解决其实我们还可以使用树状数组来动态维护一个有序序列和求解 sj 的排名,由不等式 sj-1 > si - t 和 s0 - t ~ sn - t 可以发现实际上使用的大概有 2n 个位置(s0~sn,s0 - t ~ si - t),并且 ai 的绝对值为 10 ^ 9,所以需要离散化一下;由于需要加上 si 和 si - m,所以一开始的 s0 = 0 和 s0 - m 也即 0 和 -m 需要加入到离散化数组 xs 中,枚举 1~n 将 si 和 si - m 加入到 xs 中,最后对 xs 进行去重和排序(离散化可以看成是数值映射为对应的下标)由于后面需要找到离散化对应的位置而树状数组的下标是从 1 开始的所以 xs 的元素的下标也需要从 1 开始这样才可以保持一致;并且一开始的 0 也是满足条件的所以需要将其加入到树状数组中,我们可以枚举 1~n,也即枚举所有的右端点 i, 在枚举的时候找到 sj <= si - t 的元素个数x,i - x 就是以当前的 i 为右端点满足 sj > si - t 的元素个数,然后将当前的 si 动态维护到树状数组 tr 中,枚举的时候维护 tr 的值,实际上就是树状数组对应三个函数的操作。
3. 代码如下:
python:
from typing import List
class Solution:
# lowbit, add, query为树状数组的三个经典函数
# lowbit() 函数
def lowbit(self, x: int):
return x & -x
# 树状数组中x位置加上v
def add(self, x: int, v: int, n: int, tr: List[int]):
i = x
while i < n:
tr[i] += v
i += self.lowbit(i)
# 求解区间[1, x]的区间和
def query(self, x: int, tr: List[int]):
res = 0
i = x
while i > 0:
res += tr[i]
i -= self.lowbit(i)
return res
# 获取离散化之后x在xs的位置
def get(self, x: int, count: int, xs: List[int]):
l, r = 1, count
while l < r:
mid = l + r >> 1
if xs[mid] >= x:
r = mid
else:
l = mid + 1
return r
def process(self):
n, m = map(int, input().split())
a = list(map(int, input().split()))
# s 为前缀和列表, xs 为离散化列表, tr 为树状数组列表, 最多会使用到2n+10个位置所以可以声明长度为2n+10的列表
N = 2 * n + 10
s, tr = [0] * N, [0] * N
xs = [0, -m]
for i in range(1, n + 1):
# 计算前缀和
s[i] += s[i - 1] + a[i - 1]
xs.append(s[i])
xs.append(s[i] - m)
# 去重并且将set转为列表
xs = list(set(xs))
# 排序
xs.sort()
# count为最终离散化之后的长度
count = len(xs)
# 在 xs 最前面加上0这样与下标是从1开始的与树状数组下标从1开始是对应的
xs.insert(0, 0)
# 一开始的时候0也是满足条件的所以需要将其加入到树状数组中
self.add(self.get(0, count, xs), 1, N, tr)
res = 0
for i in range(1, n + 1):
# query求解的是小于等于si - m的值所以大于si - m的元素个数为i减去query的函数值
res += i - self.query(self.get(s[i] - m, count, xs), tr)
self.add(self.get(s[i], count, xs), 1, N, tr)
return res
if __name__ == '__main__':
print(Solution().process())


