AtCoder Beginner Contest 308 E题 MEX

E题：MEX

标签：动态规划、前缀维护、MEX

题意：给定一个长度为$n$的序列$A_1,A_2,A_3,...A_n$，$A_i$的范围在$0、1、2$。再给定相同长度的字符串$B$（只含$M、E、X$三种字符）。给定规则$MEX$:指除去本身以外其他非负整数的集合中的最小值。目前要求$(B_i,B_j,B_k)$刚好是$M、E、X$，求所有对应的$MEX(A_i,A_j,A_k)$之和为多少。

题解：维护一个前缀$M$出现$0、1、2$的数量，维护一个后缀$X$出现$0、1、2$的数量，然后去枚举$E$，到每一个$E$的时候，看看对应当前的$A_i$为多少，然后考虑在这之前对应的$M$是$0、1、2$的数量，和在这之后$X$是$0、1、2$的数量。然后枚举这$9$种情况，每种情况贡献的值为：前缀$M$的数量$\ast$后缀$X$的数量$\ast$$MEX(a[i],j,k)$。$j,k$为枚举$0、1、2$的情况。

代码：

#include <bits/stdc++.h>
using namespace std;
 
typedef long long ll;
const ll N = 2e5 + 10;
ll n, a[N], m[N][10], x[N][10], ans = 0;
char b[N];
 
int cal(int x, int y, int z) {
    for (int i = 0; i <= 3; i++) {
        if (x != i && y != i && z != i) return i;
    }
    return 0;
}
 
int main() {
    cin >> n;
    for (int i = 1; i <= n; i++) cin >> a[i];
    for (int i = 1; i <= n; i++) cin >> b[i];
    for (int i = 1; i <= n; i++) {
        if (b[i] == 'M') {
            m[i][a[i]]++;
        }
        m[i][0] += m[i-1][0];
        m[i][1] += m[i-1][1];
        m[i][2] += m[i-1][2];
    }
    for (int i = n; i >= 1; i--) {
        if (b[i] == 'X') {
            x[i][a[i]]++;
        }
        x[i][0] += x[i+1][0];
        x[i][1] += x[i+1][1];
        x[i][2] += x[i+1][2];
    }
 
    for (int i = 1; i <= n; i++) {
        if (b[i] == 'E') {
            if (a[i] == 0) {
                for (int j = 0; j <= 2; j++) {
                    for (int k = 0; k <= 2; k++) {
                        ans += m[i][j] * x[i][k] * cal(0, j, k);
                    }
                }
            }
            else if (a[i] == 1) {
                for (int j = 0; j <= 2; j++) {
                    for (int k = 0; k <= 2; k++) {
                        ans += m[i][j] * x[i][k] * cal(1, j, k);
                    }
                }
            } else {
                for (int j = 0; j <= 2; j++) {
                    for (int k = 0; k <= 2; k++) {
                        ans += m[i][j] * x[i][k] * cal(2, j, k);
                    }
                }
            }
        }
    }
    cout << ans << endl;
    return 0;
}