Total number of subarrays in an array having average greater than or equal to k

Question

Given an array A of size N and an Integer k. Find the total number of subarrays with different sizes having an average greater than or equal to k.

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

Note: I have done this question by maintaining a prefix Sum array and checked for all the subarrays of all sizes. But I am getting TLE.

can we do this better(optimized one)?

Answer 1

This is pretty interesting question and a variation of a very popular question (mentioned later).

Here, we have to find number of subarrays with mean >= k.

Mathematically,

=> (a[i] + a[i + 1] + ... + a[i + l - 1]) / l >= k

=> (a[i] + a[i + 1] + ... + a[i + l - 1]) >= k * l

=> ((a[i] - k) + (a[i + 1] - k) + ... + (a[i + l - 1] - k) >= 0 ... (equation 1)

Now let's subtract all the elements of the original array by k.

After this transformation, our task is to find the number of subarrays with sum >= 0 (from equation 1)

We can break this problem in two parts:

Part 1: Finding number of subarrays with zero sum

This is pretty simple, we can do it by simply using hashmap and prefix sum. You must have solved it before.

Part 2: Finding number of subarrays with sum > k

This is pretty interesting. To solve this sub-problem, let's first build a prefix sum array for the transformed array.

Where, prefix[i] = arr[0] + arr[1] + ... + arr[i]

Now here for every index i we have to find the number of elements that are less than prefix[i] in the subarray prefix[0...(i - 1)].

Now, you can notice that this is simply a variation of the famous problem Count inversions in the array (you just need to modify that code to find number of elements greater than current to the right). This subproblem can be solved with that logic (you can go through it or you must have solved this problem earlier).

At last we have to include all the prefix sums that are > 0 (as we didn't counted it in the above parts)

The final answer would be the sum of these two counts that we get from Part 1 and Part 2. That's the whole logic behind this problem.

The time complexity would be O(n*log(n))

Here is the implementation of the logic

#include <iostream>
#include <vector>
#include <unordered_map>
#include <algorithm>

#define ll long long int

using namespace std;

ll countZeroSumSubarrays(vector<ll> &arr)
{
    unordered_map<ll, ll> mp;
    mp[0] = 1;
    ll prefix = 0, count = 0;

    for (auto &&ele : arr)
    {
        prefix += ele;

        if (mp.find(prefix) != mp.end())
            count += mp[prefix];

        mp[prefix]++;
    }

    return count;
}

void merge(ll &count, vector<ll> &v, ll left, ll mid, ll right)
{
    vector<ll> tmp(right - left + 1);
    ll i = left, j = mid + 1, k = 0;

    while (i <= mid && j <= right)
    {
        if (v[i] <= v[j])
            tmp[k++] = v[j++];
        else
        {
            count += right - j + 1;
            tmp[k++] = v[i++];
        }
    }

    while (i <= mid)
        tmp[k++] = v[i++];

    while (j <= right)
        tmp[k++] = v[j++];

    for (ll i = left; i <= right; i++)
        v[i] = tmp[i - left];
}

void mergeSort(ll &count, vector<ll> &v, ll left, ll right)
{
    if (left >= right)
        return;

    ll mid = left + (right - left) / 2;
    mergeSort(count, v, left, mid);
    mergeSort(count, v, mid + 1, right);
    merge(count, v, left, mid, right);
}

ll countSmallerElementsToRightThanSelf(vector<ll> &nums)
{
    ll N = nums.size();

    vector<ll> v(N);
    for (ll i = 0; i < N; i++)
        v[i] = nums[i];

    ll count = 0;

    mergeSort(count, v, 0, N - 1);

    return count;
}

int main(int argc, char const *argv[])
{
    ll n, k;
    cin >> n >> k;

    vector<ll> arr(n);
    for (auto &&ele : arr)
        cin >> ele;

    // Subtracting all elements by k
    for (auto &&ele : arr)
        ele -= k;

    // Creating Prefix Sum
    vector<ll> prefix(n);
    prefix[0] = arr[0];
    for (ll i = 1; i < n; i++)
        prefix[i] = prefix[i - 1] + arr[i];

    ll total = 0;
    ll subarrayWithZeroSum = countZeroSumSubarrays(arr);

    // Reversing the array so that we can find count of smaller numbers to left than self
    reverse(prefix.begin(), prefix.end());

    ll subarrayWithPositiveSum = countSmallerElementsToRightThanSelf(prefix);

    total = subarrayWithZeroSum + subarrayWithPositiveSum;
    for (auto &&ele : prefix)
        total += ele > 0;

    cout << total;
    return 0;
}

Hope it helps! If you have any doubts you can ask in the comments.