Official

E - MEX Editorial by yuto1115

解説

まずは \(j\) を全探索します (以下、\(j\) が固定されているものとします)。ただし、\(S_j=\) E を満たす \(j\) のみを考えます。\(i\)\(k\) を全探索しては間に合わないですが、この問題を解くにあたって \(i\)\(k\) そのものの値には興味がなく、\(S_i,S_k\)\(A_i,A_k\) の値にしか興味がないこと、\(S_i\)\(A_i\) の値の組は \(3\times 3 = 9\) 通りしかないことに着目します。すると、以下のようなアルゴリズムが考えられます:

  • 全ての \(a,b\ (a,b \in \lbrace 0,1,2 \rbrace)\) に対して、「\(i < j\) かつ \(S_i=\) M かつ \(A_i=a\) を満たす \(i\) の数」\(\times\)\(k> j\) かつ \(S_k=\) X かつ \(A_k=b\) を満たす \(k\) の数」\(\times\ \text{mex}(a,A_j,b)\) を答えに加算する

\(i < j\) かつ \(S_i=\) M かつ \(A_i=a\) を満たす \(i\) の数」については、「\(S_i=\) M かつ \(A_i=a\) を満たすならば \(B_i = 1\)、そうでないならば \(B_i=0\) と定めた配列 \(B\)」 の累積和を各 \(a\) についてあらかじめ求めておくことによって、すべての \(j,a\) に対する値を \(O(N)\) で求めることができます。「\(k> j\) かつ \(S_k=\) X かつ \(A_k=b\) を満たす \(k\) の数」についても同様です。よって、この問題を \(O(N)\) で解くことができました。

実装例 (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: