CF #734 (Div. 3)B2. Wonderful Coloring - 2

题目链接

Codeforces Round #734 (Div. 3)

题目大意：
用k种颜色给一个序列着色，且有以下要求
1.序列中的每个元素要么是用k种颜色绘制的，要么是没有绘制的;
2.每两个被涂成同一颜色的元素是不同的。没有两个相同颜色的值;
3.让我们计算k种颜色中每一种颜色中绘制的元素的数量-所有计算的数字必须相等;
在满足前三个条件的序列的所有着色中，该序列的着色元素总数最大。

思路：
用结构体存入序列及其下标，以序列大小排序。
然后将数值 1~k 循环赋值给结构体b，并把结构体a存的下标直接存入结构体b。
注意：这时可能不满足条件3。需要将部分赋值清零。
再将结构体b以下标从小到大排序，并输出。

#include<bits/stdc++.h>
using namespace std;
int book[301000];
struct node{
	int c,d;
}a[301000],b[301000];
bool cmp(node x,node y){//按数值排序 
	if(x.c==y.c)
	    return x.d<y.d;
	return x.c<y.c;
}
bool cmp1(node x,node y){//按下标从小到大排序 
	return x.d<y.d;
}
int main(){
	int t,n,m,i,j,k;
	scanf("%lld",&t);
	while(t--){
		memset(book,0,sizeof(book));
		scanf("%d %d",&n,&m);	    		
		for(i=1;i<=n;i++){
			scanf("%d",&a[i].c);
			a[i].d=i;
		}		  
		sort(a+1,a+n+1,cmp);
		k=0;//控制循环赋值 
		for(i=1;i<=n;i++){
			book[a[i].c]++;
			if(book[a[i].c]<=m){//使满足条件2 
				k++;
				b[i].c=k;
				if(k==m)   k=0; 
				if(k==1) j=i; //记录最后一个循环赋值的开始下标							
			}
			else
			    b[i].c=0;
			b[i].d=a[i].d;		    
		}
		if(k){//若循环赋值不完整(即不满足条件3)则需将部分赋值清零 
			for(i=j;i<=n;i++)//注意！！！ 
			    b[i].c=0;    //从最后一个循环赋值的开始下标清 
		}
		sort(b+1,b+n+1,cmp1);
		for(i=1;i<n;i++)
		    printf("%d ",b[i].c);
		printf("%d\n",b[n].c);
	}
	return 0;
}