Problem: work_life_balance_ch1 (2022 round2)

**The only differences between this chapter and [chapter 2](https://www.facebook.com/codingcompetitions/hacker-cup/2022/round-2/problems/D2) is that here, \(A_i \in \{1, 2, 3\}\), and you may swap any two 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, 2}\), or \(\mathbf{3}\)**.

You may reorder your tasks, where each *operation* lets you **swap any two elements** of \(A\).

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 100\)
\(2 \le N \le 1{,}000{,}000\)
\(1 \le M \le 1{,}000{,}000\)
\(A_i, Y_i \in \{1, 2, 3\}\)
\(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:463552492491498|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=[1, 3]\) with \(Z_2 = 1\). Whether or not we swap these tasks, they can never be equal, so \(Q_2 = -1\). The final answer is \(0 + (-1) = -1\).

The second case is depicted below:

{{PHOTO_ID:774910673835114|WIDTH:700}}

We start with \(A = [1, 1, 3, 3]\). The first update yields \(A=[1, 1, 3, 1]\). We need the sum of the first three tasks to equal the fourth. We can swap the third and fourth tasks to get \(1 + 1 + 1 = 3\), so \(Q_1 = 1\). The second update yields \(A=[3, 1, 3, 1]\). We need the sum of the first two tasks to equal the sum of the last two. This is already true, so \(Q_2 = 0\). The final answer is \(1 + 0 = 1\).

In the third case, we start with \(A = [1, 2, 3, 3, 3, 3]\). With the updates, we get:
- \(A = [1, 2, 3, 3, 3, 2], Z_1 = 3, Q_1 = 1\): we can swap \(A_1\) and \(A_6\)
- \(A = [1, 2, 3, 1, 3, 2], Z_2 = 2, Q_2 = 2\): we can swap both \(3\)s to the front
- \(A = [1, 2, 1, 1, 3, 2], Z_3 = 2, Q_3 = 1\): we can swap \(A_1\) and \(A_5\)
- \(A = [3, 2, 1, 1, 3, 2], Z_4 = 5, Q_4 = -1\): impossible

The final answer is \(1 + 2 + 1 + (-1) = 3\).


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


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


Solution: Each query depends only on the number of \(1\)s, \(2\)s, and \(3\)s present in each of the two subarrays. These counts can be computed in \(\mathcal{O}(\log N)\) time per query if three Fenwick trees, one per possible value, are maintained (in a total of \(\mathcal{O}((N+M) \log N)\) time), storing a bit for whether each index contains the value for that tree.

When it comes to evaluating each query, letting \(d\) be the difference between the subarrays' sums, our goal is to reduce \(d\) to \(0\). If we denote the subarray with the lower sum as \(L\), and the one with the higher sum as \(H\), we can influence \(d\) in the following relevant ways:

  1. Decrease \(d\) by \(4\) by swapping a \(1\) in \(L\) with a \(3\) in \(H\).
  2. Decrease \(d\) by \(2\) by swapping a \(1\) in \(L\) with a \(2\) in \(H\), or a \(2\) in \(L\) with a \(3\) in \(H\).
  3. Increase \(d\) by \(2\) by swapping a \(2\) in \(L\) with a \(1\) in \(H\), or a \(3\) in \(H\) with a \(2\) in \(H\).

Note that it's impossible to decrease \(d\) by an odd number with any swap, since if the difference between two swapped numbers is \(v\), then the sum of \(L\) will increase by \(v\) while the sum of \(H\) will decrease by \(v\), so \(d\) itself will change by \(2v\). Therefore, we can immediately report \(-1\) if \(d\) is odd.

Otherwise, we can greedily perform #1 as many times as possible without \(d\) becoming too small. One strategy is to perform #1 until either \(d < 4\), or we run out of pairs of \(1\)s and \(3\)s to swap in \(L\) and \(H\), and then perform #2 the rest of the way. Another strategy is to perform #1 until \(d\) drops to \(-2\), and then perform #3 once. We should consider both strategies (if possible) and take the one which uses fewer swaps.

To implement the two strategies, we can find the \(6\) counts of \(1\)s, \(2\)s, and \(3\)s across the two subarrays in \(\mathcal{O}(\log N)\) time with Fenwick tree queries, then simulate the swaps in \(\mathcal{O}(1)\) time.

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


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

const int INF = (int)1e9;

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

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

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

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

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

// Let L be the subarray with lower sum, and H be the one with higher sum.
// Try to reduce d = sum(H) - sum(L) down to 0 with the follow operations:
// 1. Decrease d by 4 by swapping a 1 in L with a 3 in H.
// 2. Decrease d by 2 by swapping a 1 in L with a 2 in H,
//    or a 2 in L with a 3 in H.
// 3. Increase d by 2 by swapping a 2 in L with a 1 in H,
//    or a 3 in L with a 2 in H.

// Update L/H after swapping v1 in L with v2 in H a total of swaps times.
void do_swaps(vector<int> &L, vector<int> &H, int v1, int v2, int swaps) {
  L[v1] -= swaps;
  H[v1] += swaps;
  H[v2] -= swaps;
  L[v2] += swaps;
}

// Perform op 1 until no longer possible or d = 0 or 1. Then op 2 for the rest.
int strategy_a(vector<int> L, vector<int> H, int d) {
  // We can perform op 1 at most min(num 1s in L, num 3s in H) times.
  // Each swap reduces d by 4, so we should also swap no more than floor(d/4).
  int swap13 = min(min(L[0], H[2]), d / 4);
  do_swaps(L, H, 0, 2, swap13);
  d -= 4*swap13;
  // Do the rest as needed with L[0], H[1] and L[1], H[2].
  int swap12 = min(min(L[0], H[1]), d / 2);
  do_swaps(L, H, 0, 1, swap12);
  d -= 2*swap12;
  int swap23 = min(min(L[1], H[2]), d / 2);
  do_swaps(L, H, 1, 2, swap23);
  d -= 2*swap23;
  return d != 0 ? INF : swap13 + swap12 + swap23;
}

// Perform op 1 until d becomes equal to -1, then perform op 3 once.
int strategy_b(vector<int> L, vector<int> H, int d) {
  // This time we can let d go below 0 by swapping up to ceil(d/4).
  int swap13 = min(min(L[0], H[2]), d / 4 + (int)(d % 4 != 0));
  do_swaps(L, H, 0, 2, swap13);
  d -= 4*swap13;
  // Only if d < 0, try to bring d back up to 0 with more swaps.
  int swap21 = min(min(L[1], H[0]), d < 0 ? (-d)/2 : 0);
  do_swaps(L, H, 1, 0, swap21);
  d += 2*swap21;
  int swap32 = min(min(L[2], H[1]), d < 0 ? (-d)/2 : 0);
  do_swaps(L, H, 2, 1, swap32);
  d += 2*swap32;
  return d != 0 ? INF : swap13 + swap21 + swap32;
}

long long solve() {
  int N, M;
  cin >> N >> M;
  fenwick T[]{fenwick(N), fenwick(N), fenwick(N)};
  for (int i = 1, a; i <= N; i++) {
    cin >> a;
    T[--a].set(i, 1);
  }
  long long ans = 0;
  for (int i = 0, x, y, z; i < M; i++) {
    cin >> x >> y >> z;
    --y;
    for (int v : {0, 1, 2}) {
      T[v].set(x, v == y);
    }
    vector<int> L(3), H(3);
    int sumL = 0, sumH = 0;
    for (int v : {0, 1, 2}) {
      L[v] = T[v].sum(1, z);
      H[v] = T[v].sum(z + 1, N);
      sumL += (v + 1) * L[v];
      sumH += (v + 1) * H[v];
    }
    if (sumL == sumH) {
      continue;
    }
    if (sumL > sumH) {
      swap(sumL, sumH);
      L.swap(H);
    }
    int d = sumH - sumL;
    int swaps_a = strategy_a(L, H, d), swaps_b = strategy_b(L, H, d);
    int Q = min(swaps_a, swaps_b);
    ans += (Q == INF) ? -1 : Q;
  }
  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;
}