1154E - Two Teams - CodeForces Solution

#include<bits/stdc++.h>
using namespace std;
int n,k;
vector<pair<int,int>>a;
int main(){
	cin>>n>>k;
	for(int i=1;i<=n;i++){
		int u;cin>>u;
		a.push_back({u,i});
	}
	sort(a.rbegin(),a.rend());
	set<int>s;
	for(int i=1;i<=n;i++)s.insert(i);
	queue<int>q;
	for(int i=0;i<n;i++){
		q.push(a[i].second);
	}
	int ans[n+1];
	int turn=0;
	while(s.size()){
		//cout<<1<<endl;
	vector<int>add;
	while(s.count(q.front())==0)q.pop();
	int pos=q.front();
	q.pop();
	auto it=s.find(pos);
	for(int i=0;i<=k;i++){
		add.push_back(*it);
		if(it==s.begin())break;
		it--;
	}
	it=s.find(pos);
	it++;
	for(int i=0;i<k;i++){
		if(it==s.end())break;
		else add.push_back(*it);
		it++;
	}
	for(auto x:add){
		ans[x]=turn+1;s.erase(x);
	}
	turn=turn^1;
	}
	for(int i=1;i<=n;i++)cout<<ans[i];
	return 0;
}