PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: iceknight1093
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Easy
PREREQUISITES:
None
PROBLEM:
An array is called normal if its mean and median are equal.
Given an array with elements between 1 and 2, count the number of its normal subarrays.
EXPLANATION:
The elements being between 1 and 2 is key to solving this task.
Since the mean must equal the median, and the median is an element of the array, there are really only two possibilities: the mean and median can equal 1 or 2.
Let’s look at each one in turn.
First, suppose mean = median = 1.
Now, since the elements are all \geq 1, the only way a subarray can have a mean of 1 is if all its elements are 1 - the instant a subarray contains even a single 2, its mean will be larger than 1 too.
But if all the elements are 1, the median is also surely 1, so every such subarray satisfies the conditions.
The exact same analysis applies to mean = median = 2, in that this is only possible if every element is 2.
So, what we really need to do, is count the number of subarrays whose elements are all the same.
To do this, one way is to use (very simple) dynamic programming.
Let dp_i denote the number of subarrays that end at index i, and contain all same elements.
Then,
- If A_i = A_{i-1}, we have dp_i = dp_{i-1} + 1, since any subarray ending at index i-1 can extend to include index i, and also we have one extra subarray which is the singleton [A_i].
- If A_i \neq A_{i-1}, then simply dp_i = 1.
This is easily computed in \mathcal{O}(N) time by iterating from left to right.
Finally, the answer is just the sum of all the dp_i values.
TIME COMPLEXITY:
\mathcal{O}(N) per testcase.
CODE:
Editorialist's code (PyPy3)
for _ in range(int(input())):
n = int(input())
a = list(map(int, input().split()))
ans = ct = 1
for i in range(1, n):
if a[i] != a[i-1]: ct = 0
ct += 1
ans += ct
print(ans)