求平均值≥K的子数组个数：C++数据结构选型困惑求解

约束条件 #

1 <= N <= 10^5
-10^9 <= A[i] <= 10^9

我的思路 #

设A[i]为数组前i个元素的前缀和，推导如下：

(A[j] - A[i]) / (j - i) >= K
(A[j] - A[i]) >= K * (j - i)
(A[j] - K * j) >= (A[i] - K * i) <- 记该表达式为e

由表达式e可知，若将当前索引的前缀和减去K*当前索引的值存入哈希表，只需查询其中小于等于表达式e的元素数量。我在处理第i个索引时存入哈希表的是A[i] - K * i（A[i]为数组前缀和）。

但我难以找到能在O(1)或O(logN)时间内查询小于给定元素数量的数据结构。

尝试过的数据结构 #

• 线段树：但由于约束条件过高，无法分配足够的最大空间。
• C++中的multi-set：使用upper_bound（二分查找）获取迭代器，通过计算upper_bound与begin()迭代器的差值得到小于X的元素数量，但该操作最坏时间复杂度为O(N)，导致整体时间复杂度变为O(N²)。

注：以上内容基于C语言，我使用C进行解题。希望得到相关思路与建议。

内容的提问来源于stack exchange，提问作者Ajay Kumar