AtCoder ABC 351 C - Merge the balls (4Q, 灰色, 250 点)

スタックのいい練習問題！

• 問題へのリンク

問題概要

空の列と 個のボールがある。 番目のボールの大きさは である。これから 回の操作を行う。 回目の操作では、 個目のボールを列の一番右に付け加えた後、次の手順を繰り返す。

• 列にあるボールが 1 つ以下ならば操作を終了する
• 列にあるボールのうち右から 1 番目のものと 2 番目のものの大きさが 異なるならば操作を終了する
• 列にあるボールのうち右から 1 番目のものと 2 番目のものの大きさが 等しい ならば、2 つのボールを取り除き、「取り除かれた 2 つのボールの大きさの和」の大きさのボール 1 つを列の一番右に付け加えるその後、1. に戻り、手順を繰り返す

N 回の操作の後で、列にあるボールの数を求めよ。

制約

• 

考えたこと

列はスタックで管理しよう。C++ では具体的には vector<int> 型の変数 stk で管理するとよい。右端のボールは、stk.back() によって取得できる。

一般に、大きさが のボールを 2 つ合わせると、 のボールになることに注意しよう。これに注意して、1 回の操作 (大きさ のボールを右端に入れる操作) は次のように実装できる。

while (!stk.empty() && stk.back() == A) {
    stk.pop_back();
    ++A;
}

全体として、計算量は となる。

コード

#include <bits/stdc++.h>
using namespace std;

int main() {
    int N;
    cin >> N;
    vector<int> stk;
    for (int i = 0; i < N; ++i) {
        long long A; cin >> A;
        while (!stk.empty() && stk.back() == A) {
            stk.pop_back();
            ++A;
        }
        stk.push_back(A);
    }
    cout << stk.size() << endl;
}