Official

F - Sorting Color Balls Editorial by en_translator


First of all, the minimum cost we seek equals the number of \((i,j)\) such that \(1\leq i<j\leq N\), \(C_i\neq C_j\), and \(X_i>X_j\).

We can justify in the same way as the minimum required number of swaps in bubble sort, so we just mention the outline.

Here, instead of swapping the \(i\)-th and \((i+1)\)-th balls from the left, we swap \(X_i\) and \(X_{i+1}\), and \(C_i\) and \(C_{i+1}\).

Let \(M\) be the minimum cost and \(M'\) be the number of such pairs. When the integers on the balls are non-decreasing, \(M'=0\). Swapping adjacent balls with the same color does not decrease \(M'\), while swapping them with different colors once decreases it by at most \(1\), so \(M'\leq M\) holds. All that left is to show that, when \(M'>0\) we can do some swaps of cost \(0\) and a swap of cost \(1\) to decrease \(M'\), and when \(M'=0\) we can sort to a non-decreasing sequence. In a certain state, we can first “consider consecutive balls with the same color as a block and sort them within each block in ascending order” without paying a cost. After this operation, check \(1\leq i\leq N-1\) such that the \(i\)-th and \((i+1)\)-th balls from the left has different colors; if \(X_i\leq X_{i+1}\) for all such \(i\), then the sequence is non-decreasing as a whole. Otherwise, we can swap for such \(i\) so that \(M\) is always decreased by \(1\) in exchange for a cost of \(1\). Thus, \(M'=M\) follows.

Now how can we find such \(i\)?
Let \(M_0\) be the number of \((i,j)\) such that \(1\leq i<j\leq N\) and \(X_i>X_j\),
and let \(M_k\) \((1\leq k\leq N)\) be the number of \((i,j)\) such that \(1\leq i<j\leq N\), \(C_i=C_j=k\), and \(X_i>X_j\);
then, \(M'=M_0-\displaystyle\sum_{k=1}^N M_k\).

So we just have to find the inversion number. With a data structure like a Fenwick tree, we can find \(M_0\) and \(M_k\) in an \(O(N\log N)\) and \(O(N_k\log N_k)\) time, respectively (where \(N_k\) denotes the number of balls such that \(C_i=k\)).

Noting that \(\displaystyle\sum_{k=1}^N N_k=N\), we may write the necessary number of operation with some constant \(C\) by \(CN\log N+\displaystyle\sum_{k=1}^N CN_k\log N_k \leq CN\log N+\displaystyle\sum_{k=1}^N CN_k\log N =2CN\log N\), which is bounded by \(O(N\log N)\).

Thus, we can find these value fast enough, so the problem has been solved.

Sample code in C++:

#include <bits/stdc++.h>
#include <atcoder/all>
using namespace std;
using namespace atcoder;

int main() {
	int n, x, sz;
	int c[300000];
	vector<int>a[300001];
	long long ans = 0;

	cin >> n;
	for (int i = 0; i < n; i++)cin >> c[i];
	for (int i = 0; i < n; i++) {
		cin >> x;
		a[0].push_back(x - 1);
		a[c[i]].push_back(x - 1);
	}

	fenwick_tree<long long> fw(n);
	for (int i = 0; i <= n; i++) {
		sz = a[i].size();
		for (int j = 0; j < sz; j++) {
			if (i == 0)ans += fw.sum(a[i][j] + 1, n);
			else ans -= fw.sum(a[i][j] + 1, n);
			fw.add(a[i][j], 1LL);
		}
		for (int j = 0; j < sz; j++)fw.add(a[i][j], -1LL);
	}

	cout << ans << endl;
	return 0;
}

posted:
last update: