Good Subarray Count

Question

give solution in python
D. Unique Median
time limit per test2 seconds
memory limit per test512 megabytes

An array b of m integers is called good if, when it is sorted, b_{\left\lfloor \frac{m + 1}{2} \right\rfloor} = b_{\left\lceil \frac{m + 1}{2} \right\rceil}. In other words, b is good if both of its medians are equal. In particular, \left\lfloor \frac{m + 1}{2} \right\rfloor = \left\lceil \frac{m + 1}{2} \right\rceil when m is odd, so b is guaranteed to be good if it has an odd length.

You are given an array a of n integers. Calculate the number of good subarrays^{\text{∗}} in a.

^{\text{∗}}An array x is a subarray of an array y if x can be obtained from y by the deletion of several (possibly, zero or all) elements from the beginning and several (possibly, zero or all) elements from the end.

Input
Each test contains multiple test cases. The first line contains the number of test cases t (1 \le t \le 10^4). The description of the test cases follows.

The first line of each test case contains a single integer n (1 \le n \le 10^5) — the length of the array.

The second line of each test case contains n integers a_1, a_2, \ldots, a_n (1 \le a_i \le \color{red}{\textbf{10}}) — the given array.

It is guaranteed that the sum of n over all test cases does not exceed 10^5.

Output
For each test case, output a single integer representing the number of good subarrays in a.

Example
InputCopy
3
4
1 1 1 1
5
1 10 2 3 3
10
6 3 2 3 5 3 4 2 3 5
OutputCopy
10
11
42
Note
In the first case, every subarray is good since all its elements are equal to 1.

In the second case, an example of a good subarray is b = [10, 2, 3, 3]. When it is sorted, b = [2, 3, 3, 10], so b_{\left\lfloor \frac{4 + 1}{2} \right\rfloor} = b_{\left\lceil \frac{4 + 1}{2} \right\rceil} = b_2 = b_3 = 3. Another example would be b = [1, 10, 2]. On the other hand, b = [1, 10] is not good as its two medians are 1 and 10, which are not equal.

Answer 1

Analysis

We are given an array $a$ of length $n$ , and we want to count the number of “good” subarrays. A subarray $b$ of length $m$ is called good if its two medians (the $\lfloor\frac{m+1}{2}\rfloor$ -th and $\lceil\frac{m+1}{2}\rceil$ -th elements in sorted order) are equal. Equivalently:

All odd‐length subarrays are automatically good.
- When $m$ is odd, $\lfloor \tfrac{m+1}{2}\rfloor = \lceil \tfrac{m+1}{2}\rceil$ , so the “two medians” collapse to the single middle element; they trivially match.
For even‐length subarrays of length $2k$ , being good means
$b_k \;=\; b_{k+1}$
in the sorted version of the subarray. In plainer words, the middle two elements of the sorted subarray must be the same value.

A well‐known (and somewhat surprising) key fact for this problem is:

An even‐length subarray is good if and only if it contains at least one pair of adjacent equal elements in the original array.

Let us illustrate why that works (intuitively):

If a subarray of even length contains no two adjacent equal elements at all, you can show its middle two sorted elements end up being different.
Conversely, whenever a subarray does contain some pair $a_i = a_{i+1}$ inside it, one can show that in sorted order, that value can “sit” in positions $k$ and $k+1$ .

Hence the counting boils down to:

Count all odd‐length subarrays.
Count all even‐length subarrays that contain at least one pair of adjacent equal elements.

Let’s break that down efficiently.

1) Count all subarrays and odd‐length subarrays

Total number of subarrays of an array of length $n$ is
$\frac{n \,(\,n + 1\,)}{2}.$
The number of odd‐length subarrays can be computed by noting that a subarray $[L..R]$ has odd length if and only if $L$ and $R$ have the same parity. One can show a closed‐form:
- Let $X = \bigl\lceil \tfrac{n}{2}\bigr\rceil \quad\text{(# of odd indices)}, \quad Y = \bigl\lfloor \tfrac{n}{2}\bigr\rfloor \quad\text{(# of even indices)}.$
- Then $\text{(# of odd‐length subarrays)} \;=\; \frac{X\,(\,X + 1\,)}{2} \;+\; \frac{Y\,(\,Y + 1\,)}{2}.$

Once we have the number of odd‐length subarrays, we can also get the number of even‐length subarrays by subtraction:

\text{evenSubarrays} \;=\; \frac{n\,(\,n+1\,)}{2} \;-\; \text{(# of odd‐length subarrays)}.

2) Among even‐length subarrays, count how many do not contain any adjacent‐equal pair

Equivalently, we will count how many even‐length subarrays are in a region of the array that has no two adjacent equals, i.e.\ a region that is “strictly no‐duplicates side‐by‐side.”

Scan the array from left to right and “cut” it into maximal blocks in which no two consecutive elements are equal.
Concretely:
1. Start an index start = 0.
2. For i from 0 to n-2:
  - If a[i] == a[i+1], then we end a block at index i and start a new block at i+1.
3. At the end, also close off the final block at i = n-1.
In each such block, there are no adjacent duplicates at all.
If a block has length $L$ , the number of even‐length subarrays entirely inside that block is
\underbrace{\bigl\lfloor \tfrac{L}{2}\bigr\rfloor \,\times\, \bigl\lceil \tfrac{L}{2}\bigr\rceil}_{\text{# of even‐len subarrays in a block of length }L}.
Indeed, for example:
- If $L=4$ , the even‐length subarrays inside that block are subarrays of length 2 or 4, and one can show it counts to $2\times 2=4.$
- If $L=3$ , the only even‐length subarrays are the ones of length 2, and that’s $2$ . Meanwhile, $\lfloor 3/2\rfloor \times \lceil 3/2\rceil = 1\times 2=2.$
Summing that count over all such “no‐duplicate” blocks tells us the total number of even‐length subarrays that do not contain any adjacent‐equal pair.

Denote that sum as sumBlockEven.

3) Final formula

Let

oddCount = number of odd‐length subarrays,
totalEven = number of even‐length subarrays,
sumBlockEven = number of even‐length subarrays that lie entirely in some no‐duplicate block.

Then

The number of even‐length subarrays that do contain at least one pair of adjacent equals is
$\text{evenWithDup} \;=\; \text{totalEven} \;-\; \text{sumBlockEven}.$
Finally, every odd‐length subarray is good, plus those even subarrays that have a pair of adjacent duplicates. So the total good subarrays is:
$\text{good} \;=\; \underbrace{\text{oddCount}}_{\text{all odd are good}} \;+\; \underbrace{\bigl(\text{totalEven} - \text{sumBlockEven}\bigr)}_{\text{even with a duplicate}}.$

That solves the problem in $O(n)$ time per test, which is fast enough for $n\le 10^5$ .

Python Code

python
def solve():
    import sys
    input_data = sys.stdin.read().strip().split()
    t = int(input_data[0])
    idx_data = 1
    
    # Helper to count how many odd-length subarrays in length n
    # # of odd-length subarrays = #pairs (L,R) with L,R same parity
    # If n is, say, 5 -> X=3 odd positions, Y=2 even positions
    # total odd = X*(X+1)//2 + Y*(Y+1)//2
    def count_odd_subarrays(n):
        X = (n + 1) // 2  # number of odd indices
        Y = n // 2        # number of even indices
        return X*(X+1)//2 + Y*(Y+1)//2
    
    # Number of total subarrays
    # = n*(n+1)//2
    
    # Number of even-length subarrays in length n is
    # = total_subarrays - odd_subarrays
    # = n*(n+1)//2 - count_odd_subarrays(n)
    #
    # Alternatively, you can use formula floor(n/2)*ceil(n/2).
    
    def count_even_subarrays_in_block(L):
        # For a block of length L with no adjacent duplicates,
        # the # of even-length subarrays in that block is
        # floor(L/2)*ceil(L/2)
        return (L // 2) * ((L + 1) // 2)
    
    out = []
    for _ in range(t):
        n = int(input_data[idx_data]); idx_data += 1
        arr = list(map(int, input_data[idx_data: idx_data+n]))
        idx_data += n
        
        # 1) count how many odd-length subarrays
        odd_count = count_odd_subarrays(n)
        
        # 2) count how many even-length subarrays total
        total_sub = n*(n+1)//2
        total_even = total_sub - odd_count
        
        # 3) cut the array into blocks of "no adjacent duplicates"
        blocks = []
        start = 0
        for i in range(n-1):
            if arr[i] == arr[i+1]:
                # end a block at i
                blocks.append(i - start + 1)  # length of that block
                start = i+1
        # add the final block
        blocks.append((n-1) - start + 1)
        
        # 4) sum up the even-subarray-count in each block
        sumBlockEven = 0
        for lengthL in blocks:
            sumBlockEven += count_even_subarrays_in_block(lengthL)
        
        # 5) # even subarrays that do have a pair of adj-dups
        even_with_dup = total_even - sumBlockEven
        
        # total good = all odd + all even that contain a dup
        ans = odd_count + even_with_dup
        out.append(str(ans))
    
    print('\n'.join(out))

Explanation of the Key Steps

Odd‐length subarrays are always good. We can count them in $O(1)$ from $n$ .
For even‐length subarrays, the condition “middle two elements in sorted order are equal” turns out to be equivalent to “the subarray contains (at least) one pair of adjacent equal elements somewhere inside.”
Thus we only need to subtract from the total number of even‐length subarrays those that have no adjacent duplicates at all. A subarray with no adjacent duplicates must lie entirely in some maximal block of the array that itself has no adjacent duplicates.
If a block has length $L$ (with all consecutive elements distinct), then the number of even‐length subarrays fully inside that block is $\lfloor L/2\rfloor \times \lceil L/2\rceil$ . Summing over blocks gives the total count of even‐length subarrays that do not contain a duplicate.
Subtracting that from the total # of even subarrays gives the # of even subarrays with a duplicate. Add that to the # of odd subarrays for the final answer.

That solves the problem in linear time per test and handles up to $n = 10^5$ easily.

Answer 2

Below is a fully working line-by-line version of the same solution that fixes the “no output” issue some environments face when reading everything at once with sys.stdin.read(). The core logic is unchanged; we only adjust how input is read and how output is printed. This version should pass the sample tests on Codeforces without producing an empty output.

python
def solve():
    import sys
    
    input = sys.stdin.readline  # standard fast input

    # Count how many odd-length subarrays in an array of length n
    # (# of subarrays [L..R] with R-L+1 odd)
    def count_odd_subarrays(n):
        # # of odd indices = (n+1)//2, # of even indices = n//2
        # # of odd-length subarrays = (X*(X+1) + Y*(Y+1))//2
        X = (n + 1) // 2  # number of odd indices
        Y = n // 2        # number of even indices
        return (X * (X + 1) // 2) + (Y * (Y + 1) // 2)

    # In a block of length L where all adjacent elements are distinct,
    # the number of its even-length subarrays is floor(L/2)*ceil(L/2).
    def count_even_subarrays_in_block(L):
        return (L // 2) * ((L + 1) // 2)

    t = int(input())
    answers = []
    for _ in range(t):
        n = int(input())
        arr = list(map(int, input().split()))

        # 1) Count all odd-length subarrays
        odd_count = count_odd_subarrays(n)

        # 2) Total number of subarrays
        total_sub = n * (n + 1) // 2

        # 3) Even-length subarrays = total_sub - odd_count
        total_even = total_sub - odd_count

        # 4) Find maximal blocks with no adjacent duplicates
        blocks = []
        start = 0
        for i in range(n - 1):
            if arr[i] == arr[i + 1]:
                # end a block [start..i]
                blocks.append(i - start + 1)
                start = i + 1
        # last block [start..n-1]
        blocks.append(n - start)

        # 5) Sum up how many even-length subarrays appear in these no-duplicate blocks
        sum_block_even = 0
        for lengthL in blocks:
            sum_block_even += count_even_subarrays_in_block(lengthL)

        # 6) Even subarrays that *do* contain a pair of adjacent duplicates
        even_with_dup = total_even - sum_block_even

        # 7) Final answer: all odd-length subarrays + even subarrays with a duplicate
        ans = odd_count + even_with_dup
        answers.append(str(ans))

    print("\n".join(answers))

Why this fixes the “no output” issue

In some online judges or contest environments, reading all data at once via:

python
input_data = sys.stdin.read().strip().split()

and then manually parsing can sometimes lead to index mismatches or no output if the parsing goes wrong. Switching to the standard line-by-line reading with:

python
input = sys.stdin.readline

and printing each test case’s answer immediately (or storing them to print at the end) is the most foolproof approach on Codeforces.

The core logic to count good subarrays remains the same:

Odd-length subarrays are always good.
Even-length subarrays are good if and only if they contain a pair of adjacent equal elements.
Hence we:
- Count all odd-length subarrays in $O(1)$ .
- Count how many even-length subarrays do not have duplicates, by splitting the array into blocks of pairwise-distinct elements.
- Subtract that number from total even-length subarrays to get those that do contain a duplicate.
- Sum up odd subarrays + even-with-duplicate subarrays.

With that, the sample test:

text
3
4
1 1 1 1
5
1 10 2 3 3
10
6 3 2 3 5 3 4 2 3 5

should produce:

text
10
11
42

as required.