[백준] 6198번

처음에 푼 코드

#include <bits/stdc++.h>
using namespace std;
int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);

    vector<int> v;

    int n;
    cin >> n;
    for (int i = 0; i < n; i++) {
        int in;
        cin >> in;
        v.push_back(in);
    }

    stack<int> st;
    int cnt = 0;

    for (int i = 0; i < n; i++) {
        if (st.empty())
            st.push(v[i]);
        else {
            while (!st.empty() && st.top() <= v[i]) {
                st.pop();
                cnt++;
            }
            st.push(v[i]);
        }
    }
    while (!st.empty()) {
        st.pop();
        cnt++;
    }

    cout << cnt - 1 << endl;

    return 0;
}

처음에 푼 코드가 테스트케이스는 통과했지만 로직이 이상하다는 것을 알았습니다.

로직 개선

스택에 수를 넣기 전
필요한 만큼 st.pop()을 하고
cnt += st.size()를 합니다.
그리고 st.push()를 해야 한다는 것을 알았습니다.

또한 O(n²)이므로 오버플로우가 발생할 수 있으니 cnt의 자료형을 long long으로 바꿔주었습니다.

수정된 코드

#include <bits/stdc++.h>
using namespace std;
int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);

    vector<int> v;

    int n;
    cin >> n;

    stack<int> st;
    long long cnt = 0;

    for (int i = 0; i < n; i++) {
        int in;
        cin >> in;

        if (st.empty()) {
            st.push(in);
        } else {
            while (!st.empty() && st.top() <= in) {
                st.pop();
            }
            cnt += st.size();
            st.push(in);
        }   
    }
    
    cout << cnt;

    return 0;
}