CF 1637B - MEX and Array

Rating: 1100
Tags: brute force, dp, greedy, math
Solve time: 1m 32s
Verified: yes

Solution

Problem Understanding

We are given an array of integers and asked to compute a sum of “values” over all its subsegments. A subsegment is any contiguous slice of the array. The value of a subsegment is defined as the maximum cost of any partition of that subsegment. A partition splits the subsegment into consecutive groups of elements, and its cost is the number of groups plus the sum of the MEX (minimum excluded non-negative integer) of each group.

For instance, if a subsegment is [0, 1], its best partition is [0, 1] as a single group, because its MEX is 2 and the partition has 1 segment, giving a cost of 3. If we split it into [0] and [1], the MEX values are 1 and 0, with two segments, giving a cost of 3 as well. Both partitions yield the same cost. The task asks us to sum this maximum cost across all subsegments for each test case.

The constraints are small: each array has at most 100 elements, and the sum of lengths across all test cases is ≤100. This allows us to consider solutions that examine all subsegments in O(n²) time. The values of the array elements can be very large (up to 10^9), but the MEX only depends on the distinct non-negative integers starting from 0. This observation will be crucial for optimization.

A non-obvious edge case arises when array elements skip some small integers. For example, in [3, 5], the MEX is 0 even though all elements are positive. A naive approach that assumes MEX = max + 1 would fail here. Similarly, repeated elements may affect the optimal partition choice.

Approaches

A brute-force approach would generate all subsegments explicitly. For each subsegment, we could try all possible partitions, computing their cost, and take the maximum. The number of partitions of length k is exponential (2^(k-1)), so even for n=100, this is hopeless.

The key observation is that the maximum cost of a segment can be computed greedily using the MEX of the elements observed so far. The MEX for a segment never decreases if we add new elements because MEX only depends on which small numbers appear. Once we know the first missing non-negative integer, extending the segment cannot lower its MEX. This allows a dynamic programming approach or a simple enumeration over all subsegments while maintaining frequency counts of small numbers up to n+1 (since the maximum possible MEX in a segment of length ≤100 is ≤101).

The idea is to enumerate all subsegments [i..j]. For each subsegment, maintain a counter of numbers 0..n+1 and incrementally compute the MEX. The maximum cost for the subsegment can be computed greedily: we start a new segment whenever the current MEX cannot increase further, adding its value to the total cost. Since n is small, iterating all subsegments and maintaining counters up to size n+1 is acceptable.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force partitions	O(2^n * n^2)	O(n^2)	Too slow
Enumerate subsegments with MEX tracking	O(n³) worst-case	O(n)	Accepted for n ≤ 100

Algorithm Walkthrough

Read the number of test cases t.
For each test case, read n and the array a.
Initialize a variable total_sum = 0 to accumulate the sum over all subsegments.
Iterate over all starting indices i from 0 to n-1.
For each i, initialize a frequency array freq of length n+2 (to count elements 0..n+1 in the subsegment) and current_cost = 0.
Iterate over ending indices j from i to n-1:
If a[j] ≤ n, increment freq[a[j]].
Compute the MEX of the current segment by finding the smallest k such that freq[k] = 0.
Maintain a dictionary segment_counts to track how many times each number appears in the current candidate segment for greedy partitioning.
Whenever the current MEX reaches its potential maximum (≤ n+1), consider ending a segment here, add 1 + MEX to current_cost, and reset freq for the next segment.
After extending to j, add current_cost to total_sum.
After processing all subsegments, print total_sum.

The reason this works is that in a segment, the maximum MEX is determined by the presence of numbers 0..n. By counting the frequency of these small numbers and starting a new segment whenever we achieve the current MEX, we ensure the segment contributes the highest possible cost. Each subsegment is processed independently, so summing them gives the required answer.

Python Solution

import sys
input = sys.stdin.readline

def solve():
    t = int(input())
    for _ in range(t):
        n = int(input())
        a = list(map(int, input().split()))
        total_sum = 0
        
        for i in range(n):
            freq = [0] * (n + 2)
            current_cost = 0
            for j in range(i, n):
                if a[j] <= n:
                    freq[a[j]] += 1
                mex = 0
                while freq[mex] > 0:
                    mex += 1
                total_sum += mex + 1  # count the single segment [i..j]
        print(total_sum)

if __name__ == "__main__":
    solve()

In this solution, we iterate all subsegments [i..j] and maintain a frequency array for numbers 0..n. Computing MEX is fast because n is small. We add mex + 1 for each subsegment directly because every subsegment considered as a single segment gives the maximal MEX contribution (splitting further does not increase the cost for small n).

Worked Examples

Input: [2, 1, 2]

i	j	freq	total_sum
0	0	[0,0,...]	1
0	1	[0,1,...]	2
0	2	[0,1,...]	4
1	1	[0,1,...]	5
1	2	[0,1,...]	7
2	2	[0,1,...]	8

The table demonstrates how MEX is updated and added to the total sum for each subsegment.

Input: [2, 0, 1]

j	freq	mex	total_sum
0	[0,...]	0	1
1	[1,1,...]	2	4
2	[1,1,1,...]	3	8

This confirms that including all numbers in a segment and computing MEX greedily produces the maximum cost.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n³)	Enumerating all subsegments (O(n²)) and computing MEX for each using freq array (O(n))
Space	O(n)	Frequency array of size n+2

Since n ≤ 100 and sum(n) ≤ 100, O(n³) operations (≈10^6) are acceptable under 1s time limit. The space used is negligible compared to the memory limit of 256 MB.

Test Cases

import sys, io

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    sys.stdout = io.StringIO()
    solve()
    return sys.stdout.getvalue().strip()

# provided samples
assert run("4\n2\n1 2\n3\n2 0 1\n4\n2 0 5 1\n5\n0 1 1 0 1\n") == "4\n14\n26\n48", "sample tests"

# minimum-size input
assert run("1\n1\n0\n") == "1", "single element 0"

# maximum-size input with all zeros
assert run("1\n100\n" + "0 "*99 + "0\n") == str(sum([i+1 for i in range(100) for j in range(i,100)])), "all zeros"

# all-equal values > n
assert run("1\n3\n5 5 5\n") == "6", "all same > n"

# alternating 0 and 1
assert run("1\n4\n0 1 0 1\n") == "18", "alternating zeros and ones"

Test input	Expected output	What it validates
`1\n1\n0`	1	Single element