第二回 アルゴリズム実技検定 G - ストリング・クエリ 解説

🍜

問題

問題文は本家サイトにあります:第二回 アルゴリズム実技検定 G - ストリング・クエリ

問題概要

文字列 $S$ (はじめは空文字)に対して、与えられた $Q$ 回の操作を行いなさい。なお、操作は2種類あり、 $i$ 回目の操作 $T_i$ の内容は以下の通り。

$T_i=1$ のとき:
 英小文字 $C_i$ と整数 $X_i$ が与えられる。 $S$ の末尾に $C_i$ を $X_i$ 個追加する。

$T_i=2$ のとき:
 整数 $D_i$ が与えられる。 $S$ の先頭から $D_i$ 文字を削除する。 'a', 'b', …, 'z' の各文字が削除された数をそれぞれ求め、それらを2乗し、すべてを足し合わせた結果を出力する。

制約

はじめの段階で $S$ は空文字
$C_i$ は英小文字
$1\leqq Q \leqq 10^5$
$1\leqq X_i \leqq 10^5$
$1\leqq D_i \leqq 10^5$

解説

'1 a 5' と '2 3' が入力されれば、'aaaaa' が 'aa' になるということなので、 $3^2=9$ を出力することになります。さらに '1 t 8' と '1 c 10' が入力された後に '2 21' が入力された場合を考えてみます。このとき、削除する直前の文字列 $S$ の中身は

aattttttttcccccccccc

になっています。この先頭から21文字を削除すると、a を2文字、 t を8文字、 c を10文字削除することになる(結果として、 $S$ は空文字になる)ので、\[ 2^2+8^2+10^2=168 \]を出力することになります。

やるべきことはこのような処理なのですが、問題にある通りに文字列を作るのは効率が悪いですね。長さが $10^{10}$ 程度になってしまう可能性もあります。そこで、各文字が何文字続いているかを管理することにしましょう。

先ほどの例であれば、文字列を作らずに「a が2文字、t が8文字、c が10文字続く」と管理しておけば十分だとわかります。「5文字消す」とか「13文字消す」といった処理は、前から順番に処理をしていけばいいです。こちらの方が管理しやすいです。

このように、データを蓄積した後、前から順番に処理を行う場合には、キュー queue を使うと楽になることが多いです。先ほどの例で13文字消す場合には、はじめの a を2文字消して残り11文字消す、続いて t を8文字消して残り3文字消す、というように処理をしていく、という流れです。

aが2、tが8、cが10
  ↓ 13文字消したい
  ↓ はじめのaを全部消す
tが8、cが10
  ↓ あと 11文字消したい
  ↓ はじめのtを全部消す
  :

ただ、消しきれずに文字が残ってしまう場合に困ってしまいます。c が10文字残っているところから3文字消した場合、c は残り7文字にして、以降の操作を続けていきたいですね。後ろに追加するだけなく、前に追加する処理も行いたいので、queue の代わりに、両端キュー deque を使うことにしましょう。

以上を踏まえると、C++では以下のように書くことができます。コードの下に、コードの説明が続きます。

#include <iostream>
#include <queue>
using namespace std;
using ll = long long;

int main() {
  ll Q; cin >> Q;
  deque<pair<char, ll>> q;

  for (ll i = 0; i < Q; i++) {
    ll T; cin >> T;
    if (T == 1) {
      char C; ll X;
      cin >> C >> X;
      q.emplace_back(make_pair(C, X));

    } else {
      ll D; cin >> D;
      vector<ll> cnt(26, 0);

      while (!q.empty()) {
        pair<char, ll> p = q.front();
        q.pop_front();

        if (D >= p.second) {
          cnt[p.first - 'a'] += p.second;
          D -= p.second;
        } else {
          cnt[p.first - 'a'] += D;
          p.second -= D;
          q.emplace_front(p);
          break;
        }
      }
      ll ans = 0;
      for (ll j = 0; j < 26; j++) {
        ans += (cnt[j] * cnt[j]);
      }
      cout << ans << '\\n';
    }
  }
  return 0;
}

T == 1 の場合は、文字列 C とその個数 X をペアにして、キューの後ろに追加するだけです。これで操作は終わりです。

T == 2 の場合は、D 文字削除する処理を行います。まず、キューの先頭に入っているペアを取り出します。もし、今残っている文字が全部消せるなら、全部消します。 D は今消した分だけ減らしておきます。そして、次のペアへと進みます。

もし、これから消したい文字数より、今残っている文字数の方が多い場合は、今残っている文字数を減らして、キューの先頭に戻し、削除処理を終了します。

こうして、キューが空になるか、D 文字分削除し終わるか、どちらかになれば、キューに対する処理はおしまいです。削除した文字数を各文字ごとに集計しておき、2乗の和を計算して出力します。これで操作は終わりです。