count the number of subarrays in a given array with its average being k

Question

Given an integer array a, and an integer k, we want to design an algorithm to count the number of subarrays with the average of that subarray being k. The most naive method is to traverse all possible subarrays and calculate the corresponding average. The time complexity of this naive method is O(n^2) where $n$ is the length of a. I wonder whether it is possible to do better than O(n^2).

Usually for this kind of problem, one uses prefix sum together with a hashmap, but this technique does not seem to apply here.

Answer 1

Consider a prefix sum array, call it a.

You want to find all such pairs (i, j) that (a[j]-a[i])/(j-i) == k.

Now watch the hands:

(a[j]-a[i])/(j-i) == k
a[j]-a[i] == k*(j-i)
a[j]-a[i] == k*j-k*i
a[j]-k*j  == a[i]-k*i

So if you subtract k*j from jth element of the prefix sum array, you are left with the task of counting identical pairs.