Codeforces #1486A Shifting Stacks

문제

https://codeforces.com/problemset/problem/1486/A

걸린 시간

02 : 19 : 00

풀이

C++

#include <bits/stdc++.h>
#define INF 1e9
#define all(c) c.begin(), c.end()
typedef long long ll;
typedef unsigned long long ull;
using namespace std;

int main(){
    ios::sync_with_stdio(false);
    cin.tie(NULL);
    int tc;
    cin >> tc;
    while(tc--){
        int n;
        cin >> n;
        vector<ll> h(n), s(n);
        for(int i = 0; i < n; i++){
            cin >> h[i];
            s[i] = i;
        }
        for(int i = 1; i < n; i++){
            h[i] += h[i-1];
            s[i] += s[i-1];
        }
        bool increase = true;
        for(int i = 1; i < n; i++)
            if((h[n-1-i]-s[n-1-i])+(h[n-i]-h[n-i-1]) < n-i) increase = false;
        if(increase) cout << "YES\n";
        else cout << "NO\n";
    }
    return 0;
}

n개의 스택 h 가 있고 각각 블록이 hi 개 만큼 담겨 있다. 한번에 i 번째 스택의 블록 하나를 i+1 번째 스택으로 이동시킬 수 있다고 할 때, 스택의 높이 hi 순서가 정확히 증가하도록 만들 수 있는지를 묻는 문제이다.

조건에서는 항상 낮은 인덱스에서 높은 인덱스로만 블록의 이동이 가능하다고 했다. 스택에 블록이 아무리 많아도 앞에 있는 스택으로는 블록을 이동시킬 수 없으므로 뒤로 이동시킬 수 있는 블록의 개수가 충분한지를 검사하면 된다.

문제에서 주어진 스택 배열 h 와, 조건을 만족시키는 최소 블록 개수의 스택 배열 s {0, 1, 2, 3,...} 의 구간 합(prefix sum) 을 구해두면 편리하게 구현이 가능하다.

반복문으로 제어되는 변수, 0부터 시작하는 배열 인덱스 때문에 코드가 많이 복잡해졌고, 변수를 두어 따로 관리하지 않아 정확한 인덱스 계산이 어려웠다.