Manik's blog

By Manik, history, 7 weeks ago, In English

Quite a tricky problem I was stuck on trying to solve. Would appreciate some help on how to approach it.

Given an array, compute the number of subarrays with a mean = k

1 <= n <= 1e5

1 <= a[i], k <= 1e9

Thanks!

 
 
 
 
  • Vote: I like it
  • +9
  • Vote: I do not like it

»
7 weeks ago, # |
Rev. 2   Vote: I like it +17 Vote: I do not like it

let's say dp store prefix sums.

if i<j
(dp[j]-dp[i])/(j-i)=k;
dp[j]-j*k=dp[i]-i*k;

So number of subarrays ending at j with their mean equal to k is the number of indexes who have dp[index]-index*k value equal to dp[j]-j*k.

  • »
    »
    7 weeks ago, # ^ |
      Vote: I like it 0 Vote: I do not like it

    Thank you so much for your tip, it helped me solve the problem!

    Code for anyone who is curious
  • »
    »
    7 weeks ago, # ^ |
      Vote: I like it 0 Vote: I do not like it

    Shouldn't you ask the source of the problem before you tell the approach ?

    • »
      »
      »
      7 weeks ago, # ^ |
        Vote: I like it 0 Vote: I do not like it

      That problem is very common, I've seen it several times in different platforms/contests.

»
7 weeks ago, # |
Rev. 3   Vote: I like it 0 Vote: I do not like it

Let

$$$dp[i]$$$

be total subarrays uptil index

$$$i$$$

with mean

$$$k$$$

.

So

$$$dp[i] = dp[i-1] + H[p[i] - ki]$$$

where p are the prefix sums.

Lets see how,

We want the count of all subarrays that end at index i. This is equivalent to saying, :

$$$\frac{p[i] - p[j -1]}{i-j+1} = k$$$

which is nothing but

$$$p[i] - ki = p[j-1] - k(j-1)$$$

for all

$$$j<=i$$$

we can maintain a map

$$$H$$$

and add the contribution of

$$$p[i]-ki$$$

to the answer.

»
7 weeks ago, # |
  Vote: I like it +6 Vote: I do not like it

Others have explained the solution pretty well, but I think there's a nicer way to think about it.

A neat trick that I only learned about a few days ago is called the 'average trick' (that's what I call it, anyway). If you want a subarray $$$b$$$ to have an average of $$$k$$$, then the sum of $$$(b_i - k)$$$ is $$$0$$$. So what you can do is subtract $$$k$$$ from every element. Then you only need to check that the sum of $$$b_i$$$ is $$$0$$$.

This is a more common and easier to understand problem. If you store the prefix sums, all you need to do is find the number of pairs with the same sum (because then the latter minus the former will be 0).

$$$\newline$$$

Implementation
  • »
    »
    7 weeks ago, # ^ |
      Vote: I like it 0 Vote: I do not like it

    Ahh, so when you subtract k from each element, it becomes a very standard problem (number of subarrays which sum to 0). This is incredibly elegant, thank you so much for sharing!!

  • »
    »
    7 weeks ago, # ^ |
      Vote: I like it +3 Vote: I do not like it

    This is a super cool trick, thanks for bringing it up!