AtCoder が提供しているABC（AtCoder Beginner Contest）308 のE問題をC++とPythonで解いてみました。ABC308は、2023年7月1日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 MEX（Difficulty : 1042）
問題はリンク先をご覧ください。
真ん中のEを固定して、二分探索で前後のMとXの個数を求めて問題を解きます。AtCoder Problems による Difficulty は 1042 でした。
解答案
文字列に含まれる ‘M’ と ‘E’ と ‘X’ は、0 から 2 の値を持つため、結局 9 種類のアイテムを処理することになります。この集合を M0、M1、M2、E0、E1、E2、X0、X1、X2 と表します。
真ん中の E を固定すると、その前にある Mi の個数と、その後ろにある Xi の個数から、総和に加えるべき数が分かります。
例として E が 0 のとき、以下の 9 種類の場合を考えます。
- E より前にある M0 の個数を M とするとき
- E より後ろにある X0 の個数 X に対して、M × X × 1 を加える。
- E より後ろにある X1 の個数 X に対して、M × X × 2 を加える。
- E より後ろにある X2 の個数 X に対して、M × X × 1 を加える。
- E より前にある M1 の個数を M とするとき
- E より後ろにある X0 の個数 X に対して、M × X × 2 を加える。
- E より後ろにある X1 の個数 X に対して、M × X × 2 を加える。
- E より後ろにある X2 の個数 X に対して、M × X × 3 を加える。
- E より前にある M2 の個数を M とするとき
- E より後ろにある X0 の個数 X に対して、M × X × 1 を加える。
- E より後ろにある X1 の個数 X に対して、M × X × 3 を加える。
- E より後ろにある X2 の個数 X に対して、M × X × 1 を加える。
それぞれの最後に掛けている数字は、M、E、X の MEX となります。E の前の個数と後ろの個数を求めるため、Mi と Xi の位置を格納した配列を用意して、二分探索して該当する個数を求めます。
C++ プログラム例（ABC308E）
準備として、Mi と Xi の位置を格納しておきます（33ー41行目）。Xi の個数も count_x として求めておきます（39行目）。
特定の E に対して、それより前にある Mi と Xi の個数を二分探索を使って求めます（51ー54行目）。解答となる変数 result に上記で考察したように値を加えます（55ー59行目）。下請け関数として MEX を求める mex を用意しました（6ー17行目）。
以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef unsigned long long int ull;
int mex(int m, int e, int x)
{
int result = 3;
for (int i = 0; i < 3; ++i) {
if ((m != i)&&(e != i)&&(x != i)) {
result = i;
break;
}
}
return result;
}
int main()
{
int n;
cin >> n;
vector<int> a(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
}
string s;
cin >> s;
vector<vector<int>> m(3);
vector<vector<int>> x(3);
vector<int> count_x(3);
for (int i = 0; i < n; ++i) {
if (s[i] == 'M') {
m[a[i]].push_back(i);
}
if (s[i] == 'X') {
x[a[i]].push_back(i);
++count_x[a[i]];
}
}
ull result = 0;
for (int i = 0; i < n; ++i) {
if (s[i] != 'E') {
continue;
}
// s[i] == 'E'
ull M[3];
ull X[3];
for (int j = 0; j < 3; ++j) {
M[j] = lower_bound(m[j].begin(), m[j].end(), i) - m[j].begin();
X[j] = count_x[j] - (upper_bound(x[j].begin(), x[j].end(), i) - x[j].begin());
}
for (int j = 0; j < 3; ++j) {
for (int k = 0; k < 3; ++k) {
result += M[j] * X[k] * mex(j, a[i], k);
}
}
}
cout << result << endl;
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC308E）
Python 版は C++ 版を移植しました。二分探索は bisect モジュールを使いました。関数 mex は、タプルを使うことですっきりと書けています（8行目）。
"""AtCoder Beginner Contest 308 E"""
import bisect
def mex(m, e, x):
result = 3
for i in range(3):
if i not in (m, e, x):
result = i
break
return result
n = int(input())
a = list(map(int, input().split()))
s = input()
m = [[] for _ in range(3)]
x = [[] for _ in range(3)]
count_x = [0] * 3
for i in range(n):
if s[i] == 'M':
m[a[i]].append(i)
if s[i] == 'X':
x[a[i]].append(i)
count_x[a[i]] += 1
result = 0
for i in range(n):
if s[i] != 'E':
continue
M = [0] * 3
X = [0] * 3
for j in range(3):
M[j] = bisect.bisect_left(m[j], i)
X[j] = count_x[j] - bisect.bisect_right(x[j], i)
for j in range(3):
for k in range(3):
result += M[j] * X[k] * mex(j, a[i], k)
print(result)こちらも「AC」と判定されました。
最後に
コンテスト中は、前後（’M’ と ‘X’）をあれこれ処理していて解くことができませんでした。公式解説を読み、真ん中の ‘E’ を固定すれば解けることが理解できました。
引き続き ABC の問題を紹介していきます。
参考までに真ん中の ‘E’ を固定するというヒントで最初にコーディングしたプログラムは以下でした。
#include <bits/stdc++.h>
using namespace std;
typedef unsigned long long int ull;
int main()
{
int n;
cin >> n;
vector<int> a(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
}
string s;
cin >> s;
vector<vector<int>> m(3);
vector<vector<int>> x(3);
vector<int> count_x(3);
for (int i = 0; i < n; ++i) {
if (s[i] == 'M') {
m[a[i]].push_back(i);
}
if (s[i] == 'X') {
x[a[i]].push_back(i);
++count_x[a[i]];
}
}
ull result = 0;
for (int i = 0; i < n; ++i) {
if (s[i] != 'E') {
continue;
}
// s[i] == 'E'
ull M0, M1, M2;
ull X0, X1, X2;
M0 = lower_bound(m[0].begin(), m[0].end(), i) - m[0].begin();
M1 = lower_bound(m[1].begin(), m[1].end(), i) - m[1].begin();
M2 = lower_bound(m[2].begin(), m[2].end(), i) - m[2].begin();
X0 = count_x[0] - (upper_bound(x[0].begin(), x[0].end(), i) - x[0].begin());
X1 = count_x[1] - (upper_bound(x[1].begin(), x[1].end(), i) - x[1].begin());
X2 = count_x[2] - (upper_bound(x[2].begin(), x[2].end(), i) - x[2].begin());
if (a[i] == 0) {
result += M0 * X0 * 1;
result += M0 * X1 * 2;
result += M0 * X2 * 1;
result += M1 * X0 * 2;
result += M1 * X1 * 2;
result += M1 * X2 * 3;
result += M2 * X0 * 1;
result += M2 * X1 * 3;
result += M2 * X2 * 1;
} else if (a[i] == 1) {
result += M0 * X0 * 2;
result += M0 * X1 * 2;
result += M0 * X2 * 3;
result += M1 * X0 * 2;
// result += M1 * X1 * 0;
// result += M1 * X2 * 0;
result += M2 * X0 * 3;
// result += M2 * X1 * 0;
// result += M2 * X2 * 0;
} else if (a[i] == 2) {
result += M0 * X0 * 1;
result += M0 * X1 * 3;
result += M0 * X2 * 1;
result += M1 * X0 * 3;
// result += M1 * X1 * 0;
// result += M1 * X2 * 0;
result += M2 * X0 * 1;
// result += M2 * X1 * 0;
// result += M2 * X2 * 0;
}
}
cout << result << endl;
return 0;
}このプログラムは MEX を手で計算しています。そのため冗長ですが、計算の構造はこちらのプログラムの方が分かりやすいかもしれません。AC判定もされています。
