Finding Arithmetic Mean of Subarrays Efficiently in an Array

Question

I am trying to find number of ways o calculate arithmetic mean of a subarray of an array. It boils down to this; given an array X and an integer S, how many contiguous fragments of X has arithmetic mean equal to S?

For instance, given X=[5,3,6,2] and S=4 result is 3. [5,3] , [6,2] and [5,3,6,2] has the mean 4.

X might have up to 100.000 elements. Each value of X is an integer in range of {-1.000.000.000,+1.000.000.000}. So is S. We don't round the arithmetic mean found.

My code below(on Java) works for small set of data but it is not efficient. O(n^2).

public static int returnSubsequenceCount(int[] X, int S) {
        int counter = 0;

        for (int i = 0; i < X.length; i++) {
            int[] dpSum = new int[X.length];

            dpSum[i] = X[i];

            if (X[i] == S) {
                counter++;
            }

            for (int j = i + 1; j < X.length; j++) {
                int sum = dpSum[j - 1] + X[j];

                dpSum[j] = sum;

                if ((double) sum / (j - i + 1) == S) {
                    counter++;
                }
            }
        }
        return counter;
    }

Answer 1

I will use 1-based indexing for this algorithm. This feels like one of those cases where it is OK.

Let P be the partial sums array, that is P[0] = 0 and P[i] = X[1] + ... + X[i]. Furthermore, let Q[i] = P[i] - S * i. For example,

i     0   1   2   3   4   5   6   7
-----------------------------------
X         5   3   6   2   5   5   2
P     0   5   8  14  16  21  26  28
Q     0   1   0   2   0   1   2   0

What does it mean that the average of [i,j] (including i and j) is S? With the notations above, it can be written as

(P[j] - P[i - 1]) / (j - i + 1) = S     ==>
P[j] - P[i - 1] = S * (j - i + 1)       ==>
P[j] - P[i - 1] = S * j - S * (i - 1)   ==>
P[j] - S * j = P[i - 1] - S * (i - 1)   ==>
Q[j] = Q[i - 1]

This means that any pair of equal values in Q corresponds to a range of average S. For example, the two values of 1 in Q correspond to the range [3, 6, 2, 5]. The four values of 0 in Q give rise to six ranges of average S=4: [5,3], [6,2], [5,5,2], [5,3,6,2], [6,2,5,5,2] and [5,3,6,2,5,5,2].

Therefore the following algorithm also runs in O(n log n), the same as @Polygnome's comment, but should be considerably easier to implement:

• calculate Q;
• sort Q;
• for every batch of k equal values in Q, add k * (k - 1) / 2 to the answer;
• return the answer.

This can be reduced to O(n) using a hash table or if the range of the values in Q is small enough to allow counting sort.

Answer 2

There's a trick here to obtain an O(n) algorithm:

average = (A[i] + A[i+1] ... + A[j]) / (j - i + 1)

average * (j - i + 1) = A[i] + A[i+1]...+ A[j]

Notice that since average is now multiplied by exactly the number of elements on the right side of the equation, we can subtract the average once for each one of the elements:

0 = (A[i]-average) + (A[i+1]-average) ... + (A[j]-average)

Finding contiguous sums that equal zero can be done by examining prefix sums. For each rightmost element (A[j]-average), we want to add the number of times we've seen the same prefix sum before. We make an adjustment for prefix sum 0 so as to count the full length of the array prefix if applicable:

2 1 3, avg 2

2-2 = 0    ps = 0    count = 1 (1 for the full array prefix)
1-2 = -1   ps = -1
3-2 = 1    ps = 0    count = 2 (1 for index 0 and 1 for the full array prefix)

                     total = 3