AtCoder が提供しているABC（AtCoder Beginner Contest）318 のE問題をC++とPythonで解いてみました。ABC318は、2023年9月2日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Sandwiches（Difficulty : 1004）
問題はリンク先をご覧ください。
前処理として数列にある数字毎に場所を格納してから、条件を満たす組み合わせをカウントします。AtCoder Problems による Difficulty は 1004 でした。
解答案
以下の例を考えます。1以外の数字を X、Y、Z とします（X、Y、Z はそれぞれ異なる数字を含みます）。両側が1である場合の組み合わせの個数をカウントします。
1 X X 1 Y Y Y 1 Z Z Z Z 1
Xを真ん中に持つ条件を満たす組み合わせは、$2 \times 3$ 個あります。2はXの個数を、3はXの後ろにある1の個数となります。
Yを真ん中に持つ条件を満たす組み合わせは、$2 \times 3 \times 2$ 個あります。2は、Xの前にある1の個数を、3はYの個数を、2はXの後ろのある1の個数となります。
Zを真ん中に持つ条件を満たす組み合わせは、$3 \times 4$ 個あります。3はZの前にある1の個数を、4はZの個数をとなります。
結果的に、$2 \times 3 + 2 \times 3 \times 2 + 3 \times 4$ が1に挟まれる条件を満たす組み合わせの数となります。
この操作を、各数字に対して行います。
C++ プログラム例（ABC318E）
配列 b に、それぞれの数字の場所を格納します（15－18行目）。
それぞれの数字の場所 b[i] に対して、以下の計算を行います。
- b[i] が空なら次の数字を調べる。
- i に囲まれている数字の個数を diff とする。
- diff の前にある i の個数 j + 1、後ろにある i の個数 b[i].size – j – 1 を掛けた結果を変数 result に加える。
最後に result を出力します。以下が、C＋＋プログラムとなります。
#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];
}
vector<vector<int>> b(n + 1);
for (int i = 0; i < n; ++i) {
b[a[i]].push_back(i);
}
ull result = 0;
for (int i = 1; i <= n; ++i) {
int size = b[i].size();
if (size == 0) {
continue;
}
for (int j = 0; j < size - 1; ++j) {
ull diff = b[i][j + 1] - b[i][j] - 1;
result += diff * (j + 1) * (size - j - 1);
}
}
cout << result << endl;
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC318E）
Python 版も基本的な考え方は同じです。以下となります。
"""AtCoder Beginner Contest 318 E"""
n = int(input())
a = list(map(int, input().split()))
b = [[] for i in range(n + 1)]
for i in range(n):
b[a[i]].append(i)
result = 0
for i in range(1, n + 1):
size = len(b[i])
if size == 0:
continue
for j in range(size - 1):
diff = b[i][j + 1] - b[i][j] - 1
result += diff * (j + 1) * (size - j - 1)
print(result)こちらも「AC」と判定されました。
最後に
この問題は、終了2分前に出して、WA判定でした。よく見直してみれば、中間変数（diff）を32ビット整数にしていて桁あふれしていました。今回のコンテストは、A問題とB問題で余分に時間がかかっていました。早く解いていれば、このE問題も出し直す時間があったかもしれません。コンテスト中の時間の使い方も重要だと感じました。
引き続き ABC の問題を紹介していきます。
