ABC240 D - Strange Balls C++解答例

AtCoder Beginner Contest 240 D - Strange BallsをC++で解きます。

問題

問題文をAtCoderのページより引用します。

問題文

高橋君は2以上の整数が書かれた$N$個のボールを持っており、これらを細長い筒の中に落としていきます。$i (1 \leq i \leq N)$回目には、$a_i$が書かれたボールを落とします。

ボールは特殊な材質でできており、筒の中において$k(K \geq 2)$が書かれたボールが$k$個連続すると、それら$k$個のボールは全て消えてしまいます。

各$i(1 \leq i \leq N)$について、$i$個目のボールを筒の中に落とした後、筒の中に何個のボールがあるか求めてください。

制約

• $1 \leq N \leq 2 \times 10^{5}$
• $2 \leq a_i \leq 2 \times 10^{5}$
• 入力は全て整数

解答例

シミュレーションして解く

実際に各操作とその結果をシミュレーションしていって解いてみます。

C++のstd::vectorを筒とみなすと、ボールを追加する操作はstd::vector::push_backで表現できます。
ボールを削除する操作ですが、ボールが削除されるのはボールを追加するタイミングのみであり、また削除されるボールは配列の最後の要素から数個だけです(配列の途中の要素が削除されることはない)。
そのため、ボールの削除操作はstd::vector::pop_backで表現することができます。

各ボールについて追加と削除は高々1回ずつしか発生しないので、計算量は$\mathcal{O}(N)$程度で済みます。

筒の末尾の状態を管理する

$i$回目の操作で追加するボールに書かれている数字を$k$とおきます。
ボールを追加したとき、筒の一番上にあったボールに書かれた数字が$k$と異なっていた場合は、何も起こらずボールが追加されます。
筒の一番上にあったボールに書かれた数字が$k$だったとき、以下の条件分岐が発生します。

• ボールを追加することによって、$k$が書かれたボールが筒の一番上から$k$個連続するようになったとき、筒の上から$k$個のボールを削除する
• ボールを追加しても、$k$が書かれたボールの連鎖数が$k$未満なら、何もしない

これらの計算を行うには、現在筒の一番上にあるボールに書かれている数字と、その数字が書かれたボールが筒の一番上から何個連続しているか、の2つの情報を管理しておく必要があります。
現在一番上にあるボールに書かれた数字を$T$、筒の一番上から何個連続しているかを表す数値を$C$とおきます。
これらの情報は以下のタイミングで更新されます。

• ボールが追加されたとき
  • $T \neq k$のとき、$T$を$k$で更新し、$C$を1にする
  • $T = k$のとき、$C$を1増やす
• ボールが削除されたとき
  • ボールを削除した後、筒の一番上にあるボールに書かれた数字で$T$を更新する
  • $T$と同じ数字を持つボールが筒の一番上から何個あるかを走査し、$C$に記録する

ボールを追加するたびにこの更新操作をシミュレーションすることで問題を解くことができます。

プログラム実装例

C++のプログラム実装例を以下に示します。

d.cpp

#include <iostream>
#include <vector>
#include <string>
#include <algorithm>

int main() {
  int64_t N;
  std::cin >> N;
  std::vector<int64_t> A(N);
  for (int64_t i = 0; i < N; i++) {
    std::cin >> A.at(i);
  }

  // 筒の一番上にあるボールの数字
  int64_t current_number = 0;
  // 筒の一番上にあるボールの数字の連鎖数
  int64_t current_chain = 1;
  // 筒
  std::vector<int64_t> pipe;
  for (int64_t i = 0; i < N; i++) {
    // ボールを筒の一番上に追加する
    pipe.push_back(A.at(i));

    // 筒の一番上に書かれた数字と今追加したボールの数字を比較する
    if (A.at(i) == current_number) {
      // 同じだった場合、連鎖数を1増やす
      ++current_chain;

      // 書かれた数字と連鎖数が一致したとき、ボールを削除する
      if (current_chain == current_number) {
        // 筒の一番上から連鎖数分ボールを削除する
        while (current_chain > 0) {
          pipe.pop_back();
          --current_chain;
        }

        current_chain = 0;
        // 一番上にあるボールの数字を記録する
        current_number = pipe.back();
        // 筒の一番上から走査し、連鎖数を記録する
        for (int64_t i = static_cast<int64_t>(pipe.size()) - 1; i >= 0; i--) {
          if (pipe.at(i) == current_number) {
            ++current_chain;
          } else {
            break;
          }
        }
      }
    } else {
      // 数字が異なっていた場合、各情報を更新する
      current_number = A.at(i);
      current_chain = 1;
    }

    // 出力
    std::cout << pipe.size() << std::endl;
  }

  return 0;
}

実際に提出したコードはこちら。