第3回 ドワンゴからの挑戦状 予選 C問題

問題はこちら
C: スキーリフトの相乗り - 第3回 ドワンゴからの挑戦状 予選 | AtCoder

考え方

出来る限りグループをまとめることで人数4にしてリフトを送り出したい。

保持する変数は以下の通り。

a[i] : 問題で与えられている通り、iグループの人数を入れる。
ans : 送り出すリフトの数
vector next[5]
next[i] 次の人数iのグループの場所を格納する。
cnt[i] 人数iのグループを見つけた数をカウントする。

このnext[i]とcnt[i]を使うことで、
自分が処理しているグループの次にほしい人数のグループをすぐに見つけることが可能。

#include <iostream>
#include <vector>
 
using namespace std;
 
int main(){
	int n;
	cin>>n;
	int ans=0;
	int *a;
	a = new int [n];
 
	vector<int> next[5];
	int cnt[5];
	for(int i=0;i<5;i++){
		next[i].push_back(-1);
		cnt[i] = 1;
	}
	for(int i=0;i<n;i++){
		cin>>a[i];
		next[a[i]].push_back(i);
	}
 
	for(int i=0;i<n;i++){
		int nowNum = a[i];
		if(nowNum==0){
			continue;
		}
		//cout<<i<<" "<<a[i]<<endl;
		a[i]=0;
		cnt[nowNum]++;
 
		int sum = nowNum;
		while(sum<4){
			int flag=0;
			for(int j=4-sum;j>0;j--){
				if(next[j].size()==cnt[j])continue;
 
				sum+=j;
				a[next[j].at(cnt[j])]=0;
				cnt[j]++;
				flag++;
				break;
 
			}
			if(flag==0)break;
		}
		ans++;
	}
 
	cout<<ans<<endl;
 
}