[20055号コーディネーターベルトのロボット]-C++

C++コード

#include <iostream>
#include <cstring>
#include <algorithm>
#include <cmath>
#include <queue>
#include <vector>
#include <utility>  // pair
#include <tuple>
#include <stack>
// #include <map>
#include <unordered_map>
#define ll long long
#define INF 1e9
using namespace std;

int ans = 1;
int n,k;
vector<int> A;
vector<int> robots;
int start;
int finish;

void sol() {
	while(true) {
		// no robot loaded at first
		// 1.
		rotate(A.rbegin(), A.rbegin()+1, A.rend());
		rotate(robots.rbegin(), robots.rbegin()+1, robots.rend());
		if(robots[finish]) {
			robots[finish] = 0;
		}

		// 2.
		for(int i=n-2;i>=0;--i) {
			if(robots[i]) {  // start with the oldest robot
				if(!robots[i+1] && A[i+1]) {
					robots[i+1] = 1;
					robots[i] = 0;
					A[i+1]--;
				}
			}

			if(i==n-2) {
				if(robots[finish]) {
					robots[finish] = 0;
				}
			}
		}

		// 3.
		if(A[start]) {
			robots[start] = 1;
			A[start]--;
		}

		// 4.
		int cnt = 0;
		for(int i=0;i<2*n;++i) {
			if(!A[i]) {cnt++;}
		}

		if(cnt >= k) break;
		ans++;
	}

	return;
}


int main(void) {
	// ios_base :: sync_with_stdio(false);cin.tie(NULL);cout.tie(NULL);
	
	scanf("%d%d", &n,&k);
	start = 0;
	finish = n-1;
	int num;
	for(int i=0;i<2*n;++i) {
		scanf("%d", &num);
		A.push_back(num);
		robots.push_back(0);
	}

	sol();
	
	printf("%d\n", ans);
	
	return 0;
}

実行結果