PROBLEM LINKS

Contest
Practice

Author: Constantine Sokol
Tester: Roman Rubanenko
Editorialist: Praveen Reddy Vaka

DIFFICULTY:

easy

PREREQUISITES:

Binary Indexed Tree (fenwick tree) / Segment Tree

HINT

Solve the subtask3 and try to reduce the original problem to this special case.

EXPLANATION

Let A be the given array (zero indexed).
Let sum(A,i..j) = A[i] + … + A[j]
Let avg(A,i..j) = sum(A,i..j) / (j-i+1)

Subtask1: For each pair of indices i,j (i <=j) check if average of elements between i and j is greater than or equal to K.

count=0
for i = 0 to (N-1)
    for j = i to (N-1)
        if avg(i..j) >= K
            set count = count+1
output count as answer

This solution runs in O(N^3) and hence solves only the 1st subtask.

Subtask2: You can optimize the previous solution to run in O(N^2) by making a very small observation. To compute avg(i..j) you need to compute sum(i..j). You can compute sum(A, i..j) by just adding A[j] to sum(A, 1..(j-1)). Let us look at a refined algorithm using this.

count = 0
for i = 0 to (N-1)
    sum=0
    for j = i to (N-1)
        sum = sum+A[j]
        if sum >= K*(j-i+1)
            set count = count + 1
output count as answer

Subtask3,4 and 5: Subtask3 has an additional condition of K = 0. i.e., we have to find the number of sub-arrays whose average is >= 0. It suffices to finding the number of sub-arrays whose sum >=0. We shall now reduce our problem to this special case.

Create a new array B of length same as A and set B[i] = A[i] - K for all i.

Claim: avg(A,i..j) >= K if and only if sum(B,i..j) >= 0

Proof: First we prove avg(A,i..j) >= K implies sum(B,i..j) >= 0
(A[i] +.... +A[j])/(j-i+1) >= K
A[i] +.... +A[j]) >= K*(j-i+1) = K + … + K (j-i+1times)
(A[i] - K) + … + (A[j] - K) >= 0
B[i]+...+B[j] >= 0
We can prove sum(B,i..j) >= 0 implies avg(A,i..j) >= K by just going in the backward direction

So the problem now boils down to finding the number of sub-arrays in B whose sum >= 0

Consider the PrefixSum array of B, let us call it P. (it will be of same length as B).
For each i we will have P[i] = B[0] + … + B[i]
We can compute this in O(n) time.
Note that sum(B,i..j) = P[j] - P[i-1] (imagine P[-1] to be 0)
Consider the following algorithm to compute our answer.

set count = 0
for i = 0 to (N-1)

• if P[i] >= 0 increment count by 1
• Find the number indices j such that P[j] <= P[i] and j < i, let us call the value x
• increment count by x

Output count

To find the number indices j such that P[j] <= P[i] and j < i, if we naively iterate through all indices < i our algorithm will end up being a O(n ^ 2) algorithm which is no better than our solution for subtask2. But it turns out that we can answer this particular query in O( log n) time by using an appropriate data structure.

Initially the datatstructure will not have any elements.
The idea is simple while processing a particular index i we will query the datastructure for the number of prefixSums in the data structure present already and less than or equal to current prefixSum. This has to done in O(log n) time.
We will add P[i] to the datastructure. This has to be done in O(log n) time.
Hence giving a runtime of O(n log n)

There are at least three data structures which will help you.

• Binary Indexed Tree - See setter’s and tester’s solution
• Segment Tree - See editorialist’s solution
• Balanaced Binary Search Trees - We can use an augmented balanced BST similar to Order Statistic Tree (http://en.wikipedia.org/wiki/Order_statistic_tree). But implementing a balanced BST during a contest is very tough job and hence we will skip this approach.

Refer http://community.topcoder.com/tc?module=Static&d1=tutorials&d2=binaryIndexedTrees#read to learn about Binary Indexed Trees

Refer to http://www.topcoder.com/tc?d1=tutorials&d2=lowestCommonAncestor&module=Static and link text to learn about segment trees. But these might be overkill to solve this problem. Due to lack of time I can’t explain the editorialist’s solution right now. Come back in exactly two days time and I will add the approach. Meanwhile you can have at look at the editorialist's solution and try to understand from it.

Since we have values between -10^9 and +10^9 we can’t use a BIT or a segment tree, hence we work with the positions of the values in the original array to index our data structures. Look at the solutions to see how to overcome this.

Also check out this question on quora about BIT vs Segment Tree http://www.quora.com/Data-Structures/How-does-one-decide-when-to-use-a-Segment-Tree-or-Fenwick-Tree

SOLUTIONS

Setter's Solution: SUBARR.cpp
Tester's Solution: SUBARR.cpp
Editorialist's Solution: SUBARR.java