题目描述

难度分:$1042$

输入$n(3 \leq n \leq 2 \times 10^5)$和长为$n$的数组$a(0 \leq a[i] \leq 2)$，以及长为$n$的字符串，仅包含M，E，X。

遍历所有满足$i \lt j \lt k$且$s[i]=$M且$s[j]=$E且$s[k]=$X的下标三元组$(i,j,k)$，累加 $mex(a[i],a[j],a[k])$的值，输出这个累加值。

注：$mex(a[i],a[j],a[k])$表示不在$a[i]$，$a[j]$，$a[k]$中的最小非负整数。

输入样例$1$

4
1 1 0 2
MEEX

输出样例$1$

3

输入样例$2$

3
0 0 0
XXX

输出样例$2$

0

输入样例$3$

15
1 1 2 0 0 2 0 2 0 0 0 0 0 2 2
EXMMXXXEMEXEXMM

输出样例$3$

13

算法

前后缀分解

先遍历一遍$s$，用一个数组$es$来存储所有E的位置。然后枚举所有E的位置$x$，维护$x$位置前M的个数，$x$位置后X的个数，可以开一个长度为$3$的数组$mmp$和$xmp$分别来动态维护$[1,x)$中M的数量，以及$(x,n]$中X的数量。然后分以下情况计算贡献（不考虑$Mex=0$的情况，因为对答案没有贡献）：

1. 如果中间元素$a[x]=0$，102、201这$2$种模式可以得到$Mex=3$；001、101、100这$3$种模式可以得到$Mex=2$；002、200、000、202这$4$种模式可以得到$Mex=1$。
2. 如果中间元素$a[x]=1$，不可能得到$Mex=1$，012、210这$2$种模式可以得到$Mex=3$；011、110、010这$3$种模式可以得到$Mex=2$。
3. 如果中间元素$a[x]=2$，不可能得到$Mex=2$，021、120这$2$种模式可以得到$Mex=3$；020、022、220这$3$种模式可以得到$Mex=1$。

复杂度分析

时间复杂度

遍历了两遍数组，对M和X计数是$O(1)$的，最后计算答案的时候内部计算贡献也是$O(1)$的，因此算法整体的时间复杂度为$O(n)$。

空间复杂度

开辟了一个数组$es$来记录所有E的索引，在最长情况下有$n$个E，额外空间复杂度为$O(n)$。

C++ 代码

#include <iostream>
#include <cstring>
#include <algorithm>
#include <vector>

using namespace std;
typedef long long LL;

const int N = 200010;
int n, a[N];
char s[N];

int main() {
    scanf("%d", &n);
    for(int i = 1; i <= n; i++) {
        scanf("%d", &a[i]);
    }
    scanf("%s", s + 1);
    vector<int> es;
    LL mmp[3] = {0};
    LL xmp[3] = {0};
    for(int i = 1; i <= n; i++) {
        if(s[i] == 'E') es.push_back(i);
        if(s[i] == 'X') xmp[a[i]]++;
    }
    int i = 1;
    LL ans = 0;
    for(int x: es) {
        while(i < x) {
            if(s[i] == 'M') ++mmp[a[i]];
            if(s[i] == 'X') --xmp[a[i]];
            i++;
        }
        if(xmp[0] == 0 && xmp[1] == 0 && xmp[2] == 0) break;
        if(a[x] == 0) {
            // 此时mex可以达到3
            ans += (mmp[1]*xmp[2] + mmp[2]*xmp[1])*3;                    // 102,201
            // 此时mex可以达到2
            ans += (mmp[0]*xmp[1] + mmp[1]*xmp[1] + mmp[1]*xmp[0])*2;    // 001,101,100
            // 此时mex可以达到1
            ans += mmp[0]*xmp[2] + mmp[2]*xmp[0] + mmp[0]*xmp[0] + mmp[2]*xmp[2];        // 002,200,000,202
        }else if(a[x] == 1) {
            // 此时mex可以达到3
            ans += (mmp[0]*xmp[2] + mmp[2]*xmp[0])*3;                    // 012,210
            // 此时mex可以达到2
            ans += (mmp[0]*xmp[1] + mmp[1]*xmp[0] + mmp[0]*xmp[0])*2;    // 011,110,010
        }else {
            // 此时mex可以达到3
            ans += (mmp[0]*xmp[1] + mmp[1]*xmp[0])*3;                    // 021,120
            // 此时mex可以达到1
            ans += mmp[0]*xmp[0] + mmp[0]*xmp[2] + mmp[2]*xmp[0];        // 020,022,220
        }
    }
    printf("%lld\n", ans);
    return 0;
}