👻

ABC240 D - Strange Balls C++解答例

2022/02/23に公開約2,900字

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のとき、Tkで更新し、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;
}

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

GitHubで編集を提案

Discussion

ログインするとコメントできます