AtCoder ABC 240 D - Strange Balls (3Q, 茶色, 400 点) - けんちょんの競プロ精進記録

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 240 D - Strange Balls (3Q, 茶色, 400 点)

NoviSteps3Q 茶色diff AtCoder AtCoder400点 ABC-D stack 操作をstackを用いて高速化する 入れ子構造 データ構造 データ構造テク：ペア値(要素,個数)を管理する クエリ処理問題 シミュレーション 数列

スタックを活用する系の問題！

問題概要

値 a_{1}, a_{2}, \dots, a_{N} の書かれたボールをこの順に筒に挿入していく。

挿入された際に、同じ数 k が k 個連続したら、それらがすべて消えるものとする。

i 個目のボールが挿入された瞬間の、筒内のボールの個数を逐次答えよ。

制約

  • 1 \le N \le 2 \times 10^{5}
  • 2 \le a_{i} \le 2 \times 10^{5}

考えたこと

もし仮に、「同じ数が 2 個連続したら、消える」という設定　の問題であれば、次のように、単純なスタックの活用で解くことができる。カッコ列整合性判定問題と同じ要領だ。

スタックの末尾にボールを順に挿入していく。値 k のボールを挿入するとき

  • スタックが空であるとき：スタックに値 k を挿入する
  • 空ではなく、スタックの末尾の要素値が k であるとき：その末尾要素を pop する
  • 空ではなく、スタックの末尾の要素値が k ではないとき：スタックの末尾に値 v を挿入する
  • 操作後のスタックのサイズを答える

同じ数 k が k 個連続したら消えるとき

それでは、この問題の設定に戻ろう。難しいのは

「値 k のボールを新たに挿入したときに、値 k のボールが k 個連続するかどうか」

を判定する部分だ。これを、毎回スタックを k 個遡るのでは、TLE してしまう。そこで、スタックに挿入する値を次のように修正しよう。

スタックの各要素を、ペア値 (ボールに書かれた数, その個数) とする

こうすることで、値 k のボールを新たに挿入したときに、値 k のボールが k 個連続するかどうかを判定することができる。

全体の計算量は O(N) となる。

コード

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

int main() {
    int N, res = 0;
    cin >> N;
    vector<pair<int,int>> st;
    for (int i = 0; i < N; ++i) {
        int a;
        cin >> a;
        if (st.empty() || st.back().first != a) {
            st.push_back({a,1});
            ++res;
        } else { 
            ++st.back().second;
            ++res;
            if (st.back().second == st.back().first) {
                res -= st.back().first;
                st.pop_back();
            }
        }
        cout << res << endl;
    }
}