Find the number of subarrays whose average is greater than or equal K

Question

Given an array of integers, find the number of subarrays whose average is greater than or equal to K.

Constraints:

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

My solution:

If A[i] is prefix sum upto ith index in the array then

(A[j] - A[i]) / (j - i) >= K
(A[j] - A[i]) >= K * (j - i)
(A[j] - K * j) >= (A[i] - K * i) <- Let's call this expression e

So expression e tells me If I hash the running sum till the current index along with -K * current index, then all I need to search is the number of elements less expression e.

What I was mapping in hash table after processing ith index A[i] - K * i, where A[i] is running sum of the array

But I was struggling to find a data structure which can give me number of elements less than a given element in Constant time or may be O(logN) time.

Data structures that I explored to solve this problem -

1. Segment trees but the constraints were too high for segment tree since we need to allocate the max size.
2. Multi-sets (C++) and do the upper_bound (binary search) which would give me iterator, and to get the elements less than X I can find the difference between upper_bound and begin() iterator but runtime of this could go upto O(N) and then my overall runtime goes to O(N^2).

Note: Whatever I've mentioned in the question, considering C++ as primary language since I was solving problem in C++.

Would love to hear your thoughts/suggestions.

Answer 1

You're already setting A[i] = A[i] - (K * i), so you need to find all i,j such that

A[j] - A[i] >= 0, or
A[j] >= A[i]

Assuming j>i, the number of valid pairs should just be the total number pairs minus the inversion count. You won't require any special data structure that way (an array would suffice), and inversion count calculation can be done in O(NlogN).