hackercupai/hackercup huggingface dataset:
We'll maintain two Fenwick trees representing the distribution of \(2\)s. For each index \(i\), the first Fenwick tree will store a value of \(1\) at index \(i\) iff \(A_i = 2\), while the second will store a value of \(i\) at index \(i\) iff \(A_i = 2\). These can be maintained in a total of \(\mathcal{O}((N+M) \log N)\) time.

For the \(i\)th query, we'll let \(L = A_{1..Z_i}\) and \(H = A_{(Z_i+1)..N}\). We can query the first Fenwick tree for the counts of \(2\)s in each subarray, then easily derive the sums of \(L\) and \(H\) from them. Just like in chapter 1, if the difference \(d\) of these sums is odd, they will be impossible to equalize.

For now, assume that the sum of \(L\) is less than the sum of \(H\). Equalizing them can then be thought of as "shifting" one or more \(2\)s to the left — in particular, \(d/2\) different \(2\)s will need to be moved from \(H\) into \(L\), and some \(2\)s in \(L\) may also need to be shifted left to make room for them. Firstly, if \(d/2\) plus the number of \(2\)s in \(L\) is greater than \(Z_i\), there will not be enough room to fit the the required \(2\)s. Secondly, if \(d/2\) exceeds the number of \(2\)s in \(H\), there won't be enough \(2\)s to move over. In either case, we can report \(-1\).

Otherwise, it is always possible to equalize the sums. The process will result in some suffix \(A_{s..Z_i}\) of \(L\) being entirely filled with \(2\)s, with all \(2\)s to the left of it untouched. The number of \(2\)s in that suffix will be \(c = Z_i - s + 1\). Such a suffix is valid if adding \(c\) to the number of \(2\)s in \(A_{1..(s-1)}\) comes out to the required number of \(2\)s (and if there are actually at least \(c\) \(2\)s in \(A_{s..N}\)). There may be multiple consecutive valid choices of \(s\), in which case any will do. We can find one by binary searching for \(s \in [1,Z_i]\) such that the number of \(2\)s in \(A_{s..Z_i} + d/2 = Z_i - s + 1\). This search can be done in \(\mathcal{O}(\log^2 N)\) time with the help of the first Fenwick tree.

With a valid choice of \(s\), what remains is computing the total cost of shifting the first \(c\) \(2\)s starting at index \(s\) to occupy \(A_{s..Z_i}\). We can find the index \(t\) of the last of those \(2\)s again in \(\mathcal{O}(\log^2 N)\) by binary searching with the help of the first Fenwick tree (this time for \(t \in [s, N]\) such that the number of \(2\)s in \(A_{s..t}\) equals \(c\)). Finally, we can use the second Fenwick tree to compute the sum of those \(c\) \(2\)s' original indices in \(\mathcal{O}(\log N)\) time, and subtract that sum from \(c*(s + Z_i) / 2\) (the sum of their final indices) to yield the answer \(Q_i\).

To handle the symmetric case where the sum of \(L\) is greater than \(H\), we can avoid extra thinking by maintaining two more Fenwick trees for the reversed array \(A'\). For each query, we'll run the same solution above for \(Z_i' = N - Z_i\), taking care when doing the final indices calculations.

[See David Harmeyer's solution video here.](https://youtu.be/3tuCyPdU2fo)