yutasの競技プログラミング勉強帖

競技プログラミングの問題についての解説記事を主に書いています。

E - Yet Another Sigma Problem / AtCoder Beginner Contest 353

ABC353 C++ AtCoder AtCoder Beginner Contest 500点 文字列 ハッシュ化 競技プログラミング

問題

文字列 x, \, y に対して f(x, \, y) を以下のように定義する。

  • x, \, y の最長共通接頭辞の長さを f(x, \, y) とする。

英小文字からなる n 個の文字列 ({s}_{1} , \, \cdots , \, {s}_{n}) が与えられるとき、以下の値を求めよ。

\begin{align} \sum_{i = 1}^{n - 1} \sum_{j = i + 1}^{n} f({s}_{i}, \, {s}_{j}) \end{align}

入力

まず最初の1行目に、整数 n が与えられる。

次の1行に、 n 個の文字列 {s}_{1}, \, {s}_{2}, \, \cdots , \, {s}_{n} が順に与えられる。

条件

  • 実行時間制限: 2s
  • メモリ制限: 1024MB
  • 2 \leq n \leq 3 \times {10}^{5}
  • 1 \leq | {s}_{i} |
  • | {s}_{1} | + \cdots + | {s}_{n} | \leq 3 \times {10}^{5}

出力

答えとなる値を1行で出力すること。

解法

まず、 n \leq 3 \times {10}^{5} であるので、考えられる全ての (i, \, j) の組み合わせについて f({s}_{i}, \, {s}_{j}) の値を計算する手法では、実行時間制限に間に合いません。

そこで、和を効率的に求めるために、まずは f({s}_{i}, \, {s}_{j}) の値について考えます。 f({s}_{i}, \, {s}_{j}) は「文字列 {s}_{i} と文字列 {s}_{j} の最長共通接頭辞の長さ」です。

ここで、 {s}_{i} の k 文字目を {s}_{i, \, k} と表記します。 このとき、 |{s}_{i}| 個の文字列 {s}_{i, \, 1}, \, {s}_{i, \, 1} {s}_{i, \, 2}, \, \cdots , \, {s}_{i, \, 1} \cdots {s}_{i, \, |{s}_{i}|} と、 |{s}_{j}| 個の文字列 {s}_{j, \, 1}, \, {s}_{j, \, 1}{s}_{j, \, 2}, \, \cdots , \, {s}_{j, \, 1}\cdots {s}_{j, \, |{s}_{j}|} について、文字列として一致する個数が最長共通接頭辞の長さとなります。

以上から、 i = 1, \, \cdots , \, j - 1 について、文字列 {s}_{i} の先頭から始まる連続部分文字列について、その文字列が出現する回数がカウントされているものとします。 このとき、文字列 {s}_{j} について、 i = 1, \, \cdots , \, j - 1 での f( {s}_{i}, \, {s}_{j} ) の合計は、文字列 {s}_{j} の先頭から始まる連続部分文字列について、それぞれの出現回数の合計を取ることによって求めることができます。

これを、 j = 1, \, \cdots , \, n の順に行い、各 j での合計を計算後にカウントを更新していくことで、問題の答えを求めることができます。

以上の方法では、各文字列を1文字ずつ見ていくことになるため、 O(|{s}_{1} + \cdots + |{s}_{n}|) の計算量で求めることができます。 この問題では、 |{s}_{1}| + \cdots + |{s}_{n}| \leq 3 \times {10}^{5} であるので、実行時間制限に間に合うものとして考えられます。

一方で、文字列の部分文字列の出現回数をC++での map<string, int> で管理するとします。 このとき、 {s}_{i} \leq 3 \times {10}^{5} であるので、最大で長さが 3 \times {10}^{5} の文字列を map に保存していくことになり、これではメモリ制限を超過してしまう恐れがあります。

そこで、文字列を数値として保存していくことを考えます。 具体的には、与えられる文字列は全て英小文字であるので、 a を 1、b を 2 、……というように、数値を当てはめていきます。 このとき、文字列 {s}_{i} に対して、 k 文字目を変換した値を {t}_{k} とします。 英小文字は全部で26種類あることから、部分文字列 {s}_{i, \, 1} \cdots {s}_{i, \,k} の変換値を、 \begin{align} {27}^{k - 1} {t}_{1} + {27}^{k - 2} {t}_{2} + \cdots + {t}_{k} \end{align} として定義します。 この定義により、すべての英小文字からなる文字列は、一意の値に定めることができます。

一方で、文字列は最大で 3 \times {10}^{5} の長さになることから、この値は非常に大きくなります。 これを防ぐために、ある素数 p を用意して、 \mathrm{mod} \; p の値をとることによって、値を小さくすることを考えます。

この方法を取ることで、値が大きくなることは防げますが、値が一意に定まらない可能性が出てきます。 従って、この素数 p を複数個用意して、それぞれの素数について \mathrm{mod} を取った値の配列を用意することにより、値が一意に定まらない可能性を非常に小さくすることができます。

以上の方法を実装することで、実行時間制限・メモリ制限を超過することなくこの問題の答えを求めることができます。

ソースコード

まず、上記解説内での素数 p の列を準備する関数 prepare() を以下のように実装しました。 なお、 is_prime() 関数は、引数の整数が素数の場合は true を、そうでない場合は false を返す関数です。

const int sz = 50; // 素数列のサイズ
array<ll, sz> mods; // 素数列

// 引数の整数 num が素数か否かを判定する
bool is_prime(ll num) {
  for (ll i = 2; i * i <= num; i++) {
    if (num % i == 0) {
      // もし、割り切れるようなものが存在するならば、false を返す
      return false;
    }
  }
  // forループを抜けることができたら素数なので、 true を返す
  return true;
}

// 素数列 mods を作成する関数
void prepare() {
  ll num = 5; // 現在の値として持っておく整数
  for (int i = 0; i < sz; i++) {
    // num が素数になるまで、numの値を増やし続ける
    while (!is_prime(num)) {
      num++;
    }
    mods[i] = num; // i 番目の素数として num を保存する
    num++; // num の値を増やす
  }
  return;
}

以上が実装されている状態で、答えを出力する部分を main() 関数の中に実装しました。

int n;
string s[int(3e5 + 5)];

ll ans = 0; // 答えとなる個数
const int sz = 50; // 素数列の配列サイズ
array<ll, sz> mods; // 素数列
map<array<ll, sz>, ll> cnt; // 文字列を変換した整数列の個数を保存する map

int main() {
  // 値の入力
  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> s[i];
  }

  // 素数列の準備
  prepare();

  // i = 1, …, n の順に、文字列を見ていく
  for (int i = 1; i <= n; i++) {
    // 現在の部分文字列について、整数変換した値を保存する配列
    array<ll, sz> now = {};

    // s[i] の1文字目から順番に、now の各要素をアップデートしていく
    for (int j = 0; j < s[i].length(); j++) {
      // 素数列 mods を順番に見ていき、nowの各要素をアップデートする
      for (int k = 0; k < sz; k++) {
        now[k] *= 27; // これまでの値に 27 を掛ける
        now[k] += s[i][j] - 'a' + 1; // j 文字目での値を足す
        now[k] %= mods[k]; // k 番目の素数で割った際の余りにする
      }

      ans += cnt[now]; // 現在の now について、 map に保存されている値を ans に加える
      cnt[now]++; // cnt[now] の値を 1 増やす
    }
  }
  cout << ans << endl; // 答えの出力

  return 0;
}

感想

公式解説 では、Trie木を用いた方法が紹介されていましたが、個人的には文字列をそのまま扱うのが苦手なので、このハッシュ化のような方法を使用しました。