AtCoder が提供しているABC（AtCoder Beginner Contest）353 のE問題をC++とPythonで解いてみました。ABC353は、2024年5月11日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Another Sigma Problem（Difficulty : 1217）
問題はリンク先をご覧ください。
ABC353 E問題 Another Sigma Problem
文字列を一文字毎に木の頂点に対応させて処理します。AtCoder Problems による Difficulty は 1217 でした。
解答案
前回紹介したABC273E問題（解説記事）では、木を使って問題を解きました。
同じように木を使います。木には以下の構造体を格納します。
typedef struct {
ll value;
vector<int> alpha;
} NODE;根から文字列を走査して（後で述べます）value に、この頂点を通過した文字列の数を記憶します。配列 alpha には、その後に続く文字の場所を表します。ABC273E問題では、子をたどる必要がないため、親の場所だけ格納しましたが、この問題は逆に親に戻る必要が無いため、子の場所だけ格納します。
根は、特定の文字に対応しません。
- 与えられた文字列の先頭の文字によって、子に移動します。
- もし、子が設定されていない（場所が-1）なら、子のノードを設定して、子に移動します。
- もし、子が設定されていれば、子のノードに移動します。その後、その子のノードの value を1増やして、その値を変数 result に格納します。これは、先頭の文字が一致していたそれまで処理した文字列の数を表します。
- 与えられた文字列の2文字目も同じように処理していきます。
- 全ての文字列の走査が終われば、result が求める値となります。
この木は、Trie木という名前で呼ばれているようです。この問題では、通過した文字列の数（共通接頭辞）を格納していますが、幅広い応用に使うことができそうです。
C++ プログラム例（ABC353E）
問題の制約に以下があります。
$$|S_1| + |S_2| + \cdots + |S_N| \leqq 3 \times 10^5$$
このため、木の要素数は、$3 \times 10^5$ に根を加えた数としました（19行目）。
上の考察に従い、与えられた文字列を処理します。以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
typedef struct {
ll value;
vector<int> alpha;
} NODE;
int main()
{
int n;
cin >> n;
vector<string> s(n);
for (int i = 0; i < n; ++i) {
cin >> s[i];
}
vector<NODE> a(300001);
vector<int> init_alpha(26, -1);
int max_index = 0;
ll result = 0;
a[0].value = 0;
a[0].alpha = init_alpha;
for (int i = 0; i < n; ++i) {
int now = 0;
for (int j = 0; j < s[i].length(); j++) {
int c = s[i][j] - 'a';
if (a[now].alpha[c] == -1) {
++max_index;
a[now].alpha[c] = max_index;
now = max_index;
a[now].value = 0;
a[now].alpha = init_alpha;
} else {
now = a[now].alpha[c];
++a[now].value;
result += a[now].value;
}
}
}
cout << result << endl;
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC353E）
基本的な考え方は同じです。C++の構造体の代わりにタプルを使用しました。以下がそのプログラムです。
"""AtCoder Beginner Contest 353 E"""
n = int(input())
s = list(input().split())
a = [(0, [])] * (300001)
init_alpha = [-1] * 26
max_index = 0
result = 0
a[0] = (0, init_alpha.copy())
for i in range(n):
now = 0
for _, ch in enumerate(s[i]):
c = ord(ch) - ord('a')
if a[now][1][c] == -1:
max_index += 1
a[now][1][c] = max_index
now = max_index
a[now] = (0, init_alpha.copy())
else:
now = a[now][1][c]
a[now] = (a[now][0] + 1, a[now][1])
result += a[now][0]
print(result)こちらも「AC」と判定されました。
最後に
この問題はコンテスト中に解けませんでした。コンテスト後にTrie木というデータ構造が学べました。木の応用が学べて勉強になりました。
引き続き ABC の問題を紹介していきます。
