AGC005: A - STring - pyてよn日記

問題
解法
反省
類題

　部分点解法（愚直な実装）しか実装できなかったので復習用にまとめた．基本的なデータ構造であるスタック（stack）を使って計算量を O(N²) から O(N) にするという問題．スタックを実際の問題で見たのは初めてだったので新鮮だった．

※ 追記：2019/03/03 類題に ABC120 C 問題を追加．

問題

AGC005: A - STring

f:id:pytwbf201830:20190215174445p:plain

解法

editorial

www.youtube.com

　愚直な実装では最悪計算量が O(N²) となり，部分点しかもらえない．O(N) を実現するためにはスタックを使う．以下解説から引用．

O(|X|^2)では間に合わないため、更に高速化する必要があります。実はスタックを使い、以下の操作を行うと O(|X|) で解くことが出来ます。

・文字列の文字を先頭から見ていき、
・見ている文字が S ならば、スタックに push
・見ている文字が T かつ、スタックの先頭の文字が T or 空ならばスタックに push
・見ている文字が T かつ、スタックの先頭の文字が S ならばスタックをpop
を文字列の最後まで繰り返す。この操作によりスタックに残ったものが答えの文字列です。イメージとしては、文字列を前からスキャンし、ST を発見したら適時削除していくイメージです。

スタックってこんな使い方できるんだなあとデータ構造のありがたみを感じることができた．イメージを掴むため問題のテストケースの一つを図示してみた．

f:id:pytwbf201830:20190215174550j:plain

解説通りに実装して無事解説 AC．コードは以下．

// AGC005_A
#include <iostream>
#include <string>
#include <stack>
using namespace std;

#define REP(i, n) for (ll i=0; i < n; i++)  // 0 ~ n-1
typedef long long ll;

int main() {
    // input
    string X;
    cin >> X;
    
    // calculation
    if (X.length() == 2) {
        if (X == "ST") {
            cout << 0 << "\n";
        } else {
            cout << 2 << "\n";
        }
    } else {
        // 4 <= S <= 200,000, O(N^2) は無理
        // 愚直解（部分点 200 点解法）：O(N^2)
        /*
        
        int idx = (int)X.find("ST");
        while (idx != -1) {
            X.replace(idx, 2, "");
            idx = (int)X.find("ST");
        }
        cout << X.length() << "\n";
        
        */
        // 満点解法：O(N)
        stack<char> stk;
        char ch;
        REP(i, X.length()) {
            ch = X[i];
            if (stk.empty()) {
                stk.push(ch);
            } else {
                if (stk.top() == 'S' && ch == 'T') {
                    stk.pop();
                } else {
                    stk.push(ch);
                }
            }
        }
        cout << stk.size() << "\n";
    }
    return 0;
}