2023.7.1 Atcoder 308beginner E题

2023.7.1ATcoder308(Beginner)E题

E - MEX
Problem Statement
You are given a length- $N$ sequence $A=(A_1,A_2,\dots ,A_N)$ consisting of 0, 1, and 2, and a length-$N$ string $S=S_1{S}_2\dots S_N$ consisting of M, E, and X.
Find the sum of $mex(A_i,A_j,A_k)$ over all tuples of integers $(i,j,k)$ such that $1\le i<j<k\le N$ and $S_i{S}_j{S}_k$= MEX. Here, $mex(A_i,A_j,A_k)$ denotes the minimum non-negative integer that equals neither $A_i,A_j,nor{A}_k$.
Constraints
$3\le N\le 2\times 10^5$
N is an integer.
$A_i\in 0,1,2$
$S$ is a string of length
$N$ consisting of M, E, and X.
Input
The input is given from Standard Input in the following format:
$N$
$A_1{A}_2\dots A_N$
$S$
Output
Print the answer as an integer.

大致梳理题意：给定 $n$ 个整数（0,1,或2）和 $n$ 个字符（M,E,或X），定义一个 $mex$ 函数要求满足$S_i{S}_j{S}_k$= MEX，它的值为不等于$A_i{A}_j{A}_k$的最小非负整数。然后输出所有值的和。
我们举一个样例：
4
1 1 0 2
MEEX
所有组成MEX的方案共有两个：
1.$i=1,j=2,k=4$时，$mex(A_1,A_2,A_4)$=$mex(1,1,2)$=0
2.$i=1,j=3,k=4$时，$mex(A_1,A_3,A_4)$=$mex(1,0,2)$=3
所以和为3。
大家想一想，如果只有统计ME呢？
我的想法是从前向后挨个枚举，统计当前 $M$ 的数量（当然，因为数字的不同，需要分为 $m0,m1,m2$ 来记），碰到 $E$ 就配对并且计算。
MEX同理。
从前向后挨个枚举，统计当前 $M$ 的数量（当然，因为数字的不同，需要分为 $m0,m1,m2$ 来记），碰到 $E$ 就利用乘法原理计算 $ME$ 的数量（要分为 $me00,me01,me02,me10,me11,...,me22$ 共9个，所以避免麻烦，我选择数组存储）。

#include <bits/stdc++.h>
#define int long long
using namespace std;
 
int n,a[200000],m[3],me[3][3],ans;
string s; 
 
signed main(){
	cin>>n;
	for(int i=0;i<n;++i)
		cin>>a[i];
	cin>>s;
	for(int i=0;i<n;++i){
		if(s[i]=='M')++m[a[i]];//分类统计M的数量
		else if(s[i]=='E'){
			me[0][a[i]]+=m[0];
			me[1][a[i]]+=m[1];
			me[2][a[i]]+=m[2];
		}//分类统计ME的数量
		else{
			if(a[i]==0)ans+=me[0][0]+2*me[0][1]+2*me[1][1]+3*me[1][2]+me[2][2]+me[2][0]+2*me[1][0]+3*me[2][1]+me[0][2];
			else if(a[i]==1)ans+=2*me[0][0]+2*me[0][1]+3*me[2][0]+2*me[1][0]+3*me[0][2];
			else ans+=me[0][0]+3*me[0][1]+me[2][0]+3*me[1][0]+me[0][2];
		}//计算
	}
	cout<<ans<<endl;
	return 0;
}