公式

E - MEX 解説 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;
}

投稿日時:
最終更新: