MEX and Array
Contest: Codeforces Global Round 19
Difficulty rating: 1100
Tags: brute force, dp, greedy, math

Let there be an array $$$b_1, b_2, \ldots, b_k$$$. Let there be a partition of this array into segments $$$[l_1; r_1], [l_2; r_2], \ldots, [l_c; r_c]$$$, where $$$l_1 = 1$$$, $$$r_c = k$$$, and for any $$$2 \leq i \leq c$$$ holds that $$$r_{i-1} + 1 = l_i$$$. In other words, each element of the array belongs to exactly one segment.

Let's define the cost of a partition as $$$$$$c + \sum_{i = 1}^{c} \operatorname{mex}(\{b_{l_i}, b_{l_i + 1}, \ldots, b_{r_i}\}),$$$$$$ where $$$\operatorname{mex}$$$ of a set of numbers $$$S$$$ is the smallest non-negative integer that does not occur in the set $$$S$$$. In other words, the cost of a partition is the number of segments plus the sum of MEX over all segments. Let's define the value of an array $$$b_1, b_2, \ldots, b_k$$$ as the maximum possible cost over all partitions of this array.

You are given an array $$$a$$$ of size $$$n$$$. Find the sum of values of all its subsegments.

An array $$$x$$$ is a subsegment of an array $$$y$$$ if $$$x$$$ can be obtained from $$$y$$$ by deletion of several (possibly, zero or all) elements from the beginning and several (possibly, zero or all) elements from the end.

Input format:
The input contains several test cases. The first line contains one integer $$$t$$$ ($$$1 \leq t \leq 30$$$) — the number of test cases.

The first line for each test case contains one integer $$$n$$$ ($$$1 \leq n \leq 100$$$) — the length of the array.

The second line contains a sequence of integers $$$a_1, a_2, \ldots, a_n$$$ ($$$0 \leq a_i \leq 10^9$$$) — the array elements.

It is guaranteed that the sum of the values $$$n$$$ over all test cases does not exceed $$$100$$$.

Output format:
For each test case print a single integer — the answer to the problem.

Constraints:
Time limit: 1.0 s
Memory limit: 256.0 MB

Note:
In the second test case:

- The best partition for the subsegment $$$[2, 0, 1]$$$: $$$[2], [0, 1]$$$. The cost of this partition equals to $$$2 + \operatorname{mex}(\{2\}) + \operatorname{mex}(\{0, 1\}) = 2 + 0 + 2 = 4$$$.
- The best partition for the subsegment $$$[2, 0]$$$: $$$[2], [0]$$$. The cost of this partition equals to $$$2 + \operatorname{mex}(\{2\}) + \operatorname{mex}(\{0\}) = 2 + 0 + 1 = 3$$$
- The best partition for the subsegment $$$[2]$$$: $$$[2]$$$. The cost of this partition equals to $$$1 + \operatorname{mex}(\{2\}) = 1 + 0 = 1$$$.
- The best partition for the subsegment $$$[0, 1]$$$: $$$[0, 1]$$$. The cost of this partition equals to $$$1 + \operatorname{mex}(\{0, 1\}) = 1 + 2 = 3$$$.
- The best partition for the subsegment $$$[0]$$$: $$$[0]$$$. The cost of this partition equals to $$$1 + \operatorname{mex}(\{0\}) = 1 + 1 = 2$$$.
- The best partition for the subsegment $$$[1]$$$: $$$[1]$$$. The cost of this partition equals to $$$1 + \operatorname{mex}(\{1\}) = 1 + 0 = 1$$$.

The sum of values over all subsegments equals to $$$4 + 3 + 1 + 3 + 2 + 1 = 14$$$.

Example input:
4
2
1 2
3
2 0 1
4
2 0 5 1
5
0 1 1 0 1

Example output:
4
14
26
48

Editorial:
We show, that replacing a segment of length $$$k$$$ ($$$k > 1$$$) with segments of length $$$1$$$ does not decrease the cost of the partition. Consider two cases:
1. The segment does not contain $$$0$$$.
2. The segment contains $$$0$$$.
In the first case the contribution of the segment equals to $$$1$$$ (because $$$mex = 0$$$), but the contribution of $$$k$$$ segments of length $$$1$$$ equals to $$$k$$$. So the cost increased. In the second case the contribution of the segment equals to $$$1 + mex <= 1 + k$$$, but the contribution of the segments of length $$$1$$$ would be at least $$$1 + k$$$, so the cost has not decreased.
Then it is possible to replace all segments of length more than $$$1$$$ by segments of length $$$1$$$ and not decrease the cost. So the value of the array $$$b_1, b_2, \ldots, b_k$$$ equals to $$$\sum_{i=1}^{k}{(1 + mex(\{b_i\}))}$$$ = $$$k$$$ + (the number of zeros in the array).
To calculate the total $$$value$$$ of all subsegments, you need to calculate the total length of all subsegments and the contribution of each $$$0$$$. The total length of all subsegments equals to $$$\frac{n \cdot (n + 1) \cdot (n + 2)}{6}$$$. The contribution of a zero in the position $$$i$$$ equals to $$$i \cdot (n - i + 1)$$$. This solution works in $$$O(n)$$$, but it could be implemented less efficiently.
There is also another solution, which uses dynamic programming: let $$$dp_{l, r}$$$ is the value of the array $$$a_l, a_{l + 1}, \ldots, a_r$$$. Then $$$dp_{l, r} = max(1 + mex(\{a_l, a_{l + 1}, \ldots, a_r\}), \max_{c = l}^{r - 1} (dp_{l, c} + dp_{c + 1, r}))$$$. This solution can be implemented in $$$O(n^3)$$$ or in $$$O(n^4)$$$.