题解 | #T6-小红统计区间（hard）#

提议描述非常清楚，就是求有多少个区间满足区间的和 >= k

依旧是从暴力入手：

for i in range(n):
  s = 0
  for j in range(i, n):
    s += a[j]
  if s >= k:
    cnt += 1

其中，a指的是输入数组，s指的是区间和，cnt指答案。

但是数据范围是10^5，因此暴力是无法通过的。

换个角度，区间和

如果均大于0，那么可以用滑动窗口，一旦发现一个区间大于等于k，那么右边还剩多少元素，则有多少区间也满足大于等于k的要求，具备单调性。

可是现在存在负数，因此单调性无法保证。故需要其他的办法求解

把不等式交换一下方向 一旦枚举到j，就去判断j的左边存在多少个i-1，满足该不等式。

那快速的判断一个数左边有多少小于等于它的数，引入数据结构-树状数组

算法步骤：

1、枚举右区间r，前缀和a1+a2+...+ar = s

2、利用树状数组，快速求r左边，有多少l满足 preSum[r]-pre[l-1] >= k

3、更新树状数组，因为枚举到了新的前缀和

这道题需要额外注意，值域比较大，但是n较小，因此可以对值域进行离散化。

具体代码：

n, k = map(int, input().split())
a = list(map(int, input().split()))

# 树状数组
class BIT():
    def __init__(self, n) -> None:
        self.n = n 
        self.tree = [0]*(n+1)
    
    def lowbit(self, i):
        return i & -i 

    def add(self, i, x):
        while i < self.n:
            self.tree[i] += x 
            i += self.lowbit(i)
    
    def query(self, i):
        res = 0
        while i:
            res += self.tree[i]
            i -= self.lowbit(i)
        return res 
            

# 离散化
index = {} # index存放离散化之后的下标
s = 0 # 前缀和s
brr = set()

for i in range(n):
    s += a[i]
    brr.add(s)
    brr.add(s-k)
    
brr.add(0)
brr = list(brr)
brr.sort()

for i in range(len(brr)):
    index[brr[i]] = i + 1

# 查询

s = 0
res = 0
tr = BIT(len(brr))
tr.add(index[0], 1)

for i,x in enumerate(a):
    s += x 
    res += tr.query(index[s-k])
    tr.add(index[s], 1)

print(res)

我自己也录制了一下视频讲解，尝试成为一个小up，哈哈哈哈

牛客周赛 Round 28 - F（离散化+树状数组）