E - MEX Editorial by en_translator
We exhaustively try all \(j\) such that \(S_j=\) E
(hereinafter we take a fixed \(j\)). While we cannot enumerate all \(i\) and \(k\), notice that we are not interested in the value of \(i\) and \(k\) itself, but only in \(S_i\), \(S_k\), \(A_i\) and \(A_k\), and that there are \(3\times 3 = 9\) kinds of value for the pair of values of \(S_i\) and \(A_i\). Then we come up with the following algorithm:
- For each \(a,b\ (a,b \in \lbrace 1,2,3 \rbrace)\), add the following value to the answer: “the number of \(i\)s with \(i < j\), \(S_i=\)
M
, and \(A_i=a\)” \(\times\) “the number of \(k\)s with \(k> j\), \(S_k=\)X
, and \(A_i=b\)” \(\times\ \text{mex}(a,A_j,b)\).
“The number of \(i\)s with \(i < j\), \(S_i=\) M
, and \(A_i=a\)” can be found for all \(j\) and \(a\) in a total of \(O(N)\) time by precomputing for each \(a\) the cumulative sums of the array \(B\) where \(B_i = 1\) if \(S_i=\) M
and \(A_i=a\) and \(B_i=0\) otherwise. Same applies to “the number of \(k\)s with \(k> j\), \(S_k=\) X
, and \(A_i=b\).” Thus, the problem has been solved in a total of \(O(N)\) time.
Sample code (C++):
#include<bits/stdc++.h>
using namespace std;
using ll = long long;
int mex(int x, int y, int z) {
for (int i = 0; i < 3; i++) {
if (x != i and y != i and z != i) return i;
}
return 3;
}
int main() {
int n;
cin >> n;
vector<int> a(n);
for (int &i: a) cin >> i;
string s;
cin >> s;
vector cnt_l(n + 1, vector<int>(3));
vector cnt_r(n + 1, vector<int>(3));
for (int i = 0; i < n; i++) {
cnt_l[i + 1] = cnt_l[i];
if (s[i] == 'M') ++cnt_l[i + 1][a[i]];
}
for (int i = n - 1; i >= 0; i--) {
cnt_r[i] = cnt_r[i + 1];
if (s[i] == 'X') ++cnt_r[i][a[i]];
}
ll ans = 0;
for (int i = 0; i < n; i++) {
if (s[i] != 'E') continue;
for (int j = 0; j < 3; j++) {
for (int k = 0; k < 3; k++) {
ans += (ll) cnt_l[i][j] * cnt_r[i + 1][k] * mex(j, a[i], k);
}
}
}
cout << ans << endl;
}
posted:
last update: