Where's My Water?

Row Components

At row $y$, column $i$ contains water iff $a_i < y$. So the water cells in that row are split into maximal contiguous intervals.

If a drain's column lies inside one such interval, then every water tile in that interval can move horizontally to the drain's column, then move down to the drain. So at each row, a drain collects exactly the whole component containing its column.

Also, if we choose column $i$, putting the drain higher is strictly worse. Move it down to row $a_i+1$: every tile that could reach the old drain can continue downward in the same column. So each drain is just a chosen column.

The Hidden Tree

As $y$ increases, more columns become water, so components only merge. This gives a laminar interval structure.

Consider some interval $[l,r]$ that is known to be surrounded by barriers until height $U$. Let

$$m = \max(a_l, a_{l+1}, \dots, a_r).$$

For every row $y$ with $m < y \le U$, all columns in $[l,r]$ are water, so $[l,r]$ is one component. That component contributes

$$w = (U-m)(r-l+1)$$

water tiles if at least one selected drain column lies inside $[l,r]$.

For rows $y \le m$, every column with height $m$ is still dirt, so those columns split $[l,r]$ into independent smaller intervals. We recurse on the maximal segments between positions where $a_i=m$, using upper height $m$.

This is basically a max-Cartesian-tree decomposition, but we do not need to explicitly build the tree. Just recursively split by the maximum. Dirt does the divide-and-conquer for us. Very considerate dirt.

DP

Let $solve(l,r,U)$ return $dp[0], dp[1], dp[2]$, where $dp[k]$ is the maximum water collected inside this recursive interval using at most $k$ drains placed in $[l,r]$.

Inside $[l,r]$:

1. Find $m=\max(a_l,\dots,a_r)$.
2. The current component weight is $w=(U-m)(r-l+1)$.
3. Recursively solve every child segment between columns where $a_i=m$.
4. Combine children with a capacity-$2$ knapsack:

$$best[c] = \max \text{ water from processed children using at most } c \text{ drains}.$$

Now the key transition is simple:

$$dp[0]=0,$$

and for $k \in \{1,2\}$,

$$dp[k] = w + best[k].$$

Why can we always add $w$ when $k \ge 1$? Because either some drain is already placed in a child, hence it lies inside $[l,r]$, or we can place one drain on a maximum-height column of this interval to collect the current component. If $w=0$, adding it is harmless anyway.

The final answer is $solve(1,n,h)[2]$.

Correctness

Every water tile belongs to exactly one row-component at its row. Our recursive decomposition assigns that component to exactly one interval node: the node whose height band contains that row.

A selected drain column collects a node's whole component iff the selected column lies inside that node's interval. Therefore, maximizing drained water is exactly the same as choosing up to two columns to maximize the total weight of interval nodes containing at least one chosen column.

The children of any interval are disjoint, so once the current component weight is handled, the remaining optimization is just distributing $0$, $1$, or $2$ drains among independent child intervals. The knapsack transition checks all such distributions, so it is exhaustive. Since the current node only cares whether at least one drain exists inside it, adding $w$ for $k\ge 1$ is exactly right.

Thus the DP computes the optimal value for every recursive interval, and the root gives the optimal answer for the whole grid.

Complexity

Each recursive call scans its interval to find the maximum and split children. In the worst case, the total scanned length is $O(n^2)$. The DP capacity is only $2$, so child merging is constant-factor tiny.

Across all test cases, $\sum n \le 2000$, so this easily fits.

$$\text{Time} = O(n^2), \qquad \text{Memory} = O(n)$$