Problem: work_life_balance_ch2 (2022 round2)

**The only differences between this chapter and [chapter 1](https://www.facebook.com/codingcompetitions/hacker-cup/2022/round-2/problems/D1) is that here, \(A_i \in \{1, 2\}\), and you may only swap adjacent elements of \(A_i\) with a single operation.**

As a Metal Platforms employee, you place a high value on your work-life balance. Boss Rob has assigned you \(N\) tasks, the \(i\)th of which takes \(A_i\) minutes to finish, where \(A_i\) is either **\(\mathbf{1}\) or \(\mathbf{2}\)**.

You may reorder your tasks, where each *operation* lets you **swap any two adjacent elements** of \(A\) (\(A_i\) and \(A_j\) for some \(i\) and \(j\) such that \(|i - j| = 1\)).

To reflect how often task requirements change in the real world, there will be \(M\) updates made to the task completion times, with the \(i\)th update setting \(A_{X_i}\) to \(Y_i\).

After completing the \(i\)th update, you would like to know if it's possible to balance the time spent at work versus at home, namely if you hope to finish the first \(Z_i\) tasks at work and the rest at home. Specifically, let \(Q_i\) be the minimum number of swap operations which must be theoretically performed so that \(A_1 + ... + A_{Z_i} = A_{Z_i + 1} + ... + A_{N}\), with \(Q_i = -1\) if it's impossible. Note that it's possible for \(Q_i\) to be \(0\), if the subarrays already have equal sums.

To reduce the size of the output, please compute the sum of \(Q_1\), ..., \(Q_M\).


# Constraints

\(1 \le T \le 95\)
\(2 \le N \le 1{,}000{,}000\)
\(1 \le M \le 1{,}000{,}000\)
\(A_i, Y_i \in \{1, 2\}\)
\(1 \le X_i \le N\)
\(1 \le Z_i \le N-1\)

The sum of \(N+M\) across all test cases is at most \(7{,}000{,}000\).


# Input Format

Input begins with a single integer \(T\), the number of test cases. For each test case, there is first a line containing two space-separated integers \(N\) and \(M\). Then, there is a line containing \(N\) digits, \(A_1\), \(A_2\), ..., \(A_N\). Then, \(M\) lines follow, the \(i\)th of which contains three space-separated integers \(X_i\), \(Y_i\), and \(Z_i\).


# Output Format

For the \(i\)th test case, output `"Case #i: "` followed by a single integer, \(Q_1 + ... + Q_M\).


# Sample Explanation

The first case is depicted below:

{{PHOTO_ID:2337609416394611|WIDTH:700}}

We start with \(A = [1, 2]\). The first update yields \(A=[1, 1]\) with \(Z_1 = 1\). The first and second tasks are already equal, so \(Q_1 = 0\) as no swaps are necessary. The second update yields \(A=[2, 1]\), and after the third update we still have \([2, 1]\). In neither case can we equally split the tasks, so \(Q_2 = Q_3 = -1\). The final answer is \(0 + (-1) + (-1) = -2\).

The second case is depicted below:

{{PHOTO_ID:402660235386017|WIDTH:700}}

We start with \(A = [1, 1, 1, 2]\). The first update yields \(A=[2, 1, 1, 2]\) with \(Z_1 = 2\). The two subarrays both sum to \(3\) already, so \(Q_1 = 0\). The second update yields \(A=[2, 2, 1, 2]\). It's impossible to split the tasks equally, so \(Q_2 = -1\). The third update yields \(A=[2, 2, 1, 1]\) with \(Z_3 = 2\). We can swap the second and third tasks so that both halves of the array sum to \(3\), so \(Q_3 = 1\). The final answer is \(0 + (-1) + 1 = 0\).

In the third case, we start with \(A = [1, 1, 1, 1, 2, 2, 2, 2]\). With the updates, we get:
- \(A = [1, 1, 1, 1, 2, 2, 2, 2], Z_1 = 4\) and \(Q_1 = 4\): move the left two \(2\)s two spaces left
- \(A = [1, 1, 1, 1, 2, 2, 2, 2], Z_2 = 3\) and \(Q_2 = 12\): move the left three \(2\)s four spaces left
- \(A = [1, 1, 1, 1, 2, 2, 2, 2], Z_3 = 5\) and \(Q_3 = 0\)
- \(A = [2, 1, 1, 1, 2, 2, 2, 2], Z_4 = 4\) and \(Q_4 = -1\)
- \(A = [2, 1, 2, 1, 2, 2, 2, 2], Z_5 = 4\) and \(Q_5 = 1\): swap \(A_4\) and \(A_5\)

The final answer is \(4 + 12 + 0 + (-1) + 1 = 16\).


Sample input:
3
2 3
1 2
2 1 1
1 2 1
1 2 1
4 3
1 1 1 2
1 2 2
2 2 2
4 1 2
8 5
1 1 1 1 2 2 2 2
5 2 4
7 2 3
6 2 5
1 2 4
3 2 4


Sample output:
Case #1: -2
Case #2: 0
Case #3: 16


Solution: 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)


Code: #include <cassert>
#include <iostream>
#include <vector>
using namespace std;

using LL = long long;

struct fenwick {  // For 1-based indices.
  vector<LL> a, t;

  fenwick(int N) : a(N + 1), t(N + 1) {}

  LL sum(int hi) const {
    LL res = 0;
    for (; hi > 0; hi -= hi & -hi) {
      res += t[hi];
    }
    return res;
  }

  void inc(int i, LL x) {
    a[i] += x;
    for (; i < (int)t.size(); i += i & -i) {
      t[i] += x;
    }
  }

  LL at(int i) const { return a[i]; }
  LL sum(int lo, int hi) const { return sum(hi) - sum(lo - 1); }
  void set(int i, LL x) { inc(i, x - a[i]); }
};

template <typename Predicate>
int binary_search_last_true(int lo_incl, int hi_excl, Predicate pred) {
  int mid, lo = lo_incl, hi = hi_excl - 1;
  while (lo < hi) {
    mid = lo + (hi - lo + 1) / 2;
    if (pred(mid)) {
      lo = mid;
    } else {
      hi = mid - 1;
    }
  }
  return pred(lo) ? lo : hi_excl;
}

LL solve() {
  int N, M;
  cin >> N >> M;
  auto rev = [=](int i) -> int { return N - i + 1; };
  fenwick Tv_fwd(N), Ti_fwd(N), Tv_rev(N), Ti_rev(N);
  for (int i = 1, a; i <= N; i++) {
    cin >> a;
    Tv_fwd.set(i, a == 2);
    Ti_fwd.set(i, a == 2 ? i : 0);
    Tv_rev.set(rev(i), a == 2);
    Ti_rev.set(rev(i), a == 2 ? i : 0);
  }
  LL ans = 0;
  for (int i = 0, x, y, z; i < M; i++) {
    cin >> x >> y >> z;
    Tv_fwd.set(x, y == 2);
    Ti_fwd.set(x, y == 2 ? x : 0);
    Tv_rev.set(rev(x), y == 2);
    Ti_rev.set(rev(x), y == 2 ? x : 0);
    int L2cnt = Tv_fwd.sum(1, z), H2cnt = Tv_fwd.sum(z + 1, N);
    int sumL = 1 * (z - L2cnt) + 2 * L2cnt;
    int sumH = 1 * (N - z - H2cnt) + 2 * H2cnt;
    if (sumL == sumH) {
      continue;
    }
    auto getQ = [&](
      const fenwick &Tv, const fenwick &Ti, int d, int z, bool is_rev
    ) -> LL {
      int req2s = d / 2, tgt2s = Tv.sum(1, z) + req2s;
      if (
        d % 2 == 1 ||  // Can't reduce d by an odd number with swaps.
        tgt2s > z ||  // Not enough room to fit all the 2's to make d = 0.
        req2s > Tv.sum(z + 1, N)  // Not enough 2's from H to move over to L.
      ) {
        return -1;
      }
      // Assuming sumL < sumH, check that a choice of s will make sumL >= sumH
      // if we fill a A[s..z] by shifting 2's from A[(z+1)..N] leftwards, all
      // without touching the 2's in A[1..(s-1)].
      int s = binary_search_last_true(1, z + 1, [&](int s) {
        // On the domain s=1..z, this predicate yields a pattern 1111000, as
        // we'll need force sufficiently many 2's in A[s..z] to hit the target
        // number of 2s required to make sumL = sumH.
        return Tv.sum(1, s - 1) + (z - s + 1) >= tgt2s;
      });
      assert(s <= z);
      int c = z - s + 1;
      // Now compute the total cost of shifting the 2s from the range s..t
      // leftwards to fill the indices s..z.
      int t = 1 + binary_search_last_true(s, N + 1, [&](int t) {
        // On the domain t=s..N, this predicate yields a pattern 1111000, as
        // we'll need to choose sufficiently large t for there to be enough 2s.
        return Tv.sum(s, t) < c;
      });
      assert(t <= N);
      // Q_i = (sum of final indices of 2s) - (sum of initial indices of 2s).
      if (is_rev) {
        return (LL)c*(rev(s) + rev(z))/2 - Ti.sum(s, t);
      }
      return Ti.sum(s, t) - (LL)c*(s + z)/2;
    };
    if (sumL < sumH) {
      ans += getQ(Tv_fwd, Ti_fwd, sumH - sumL, z, false);
    } else {
      ans += getQ(Tv_rev, Ti_rev, sumL - sumH, rev(z) - 1, true);
    }
  }
  return ans;
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);
  int T;
  cin >> T;
  for (int t = 1; t <= T; t++) {
    cout << "Case #" << t << ": " << solve() << endl;
  }
  return 0;
}