晚间测试10 A. 分组 带权并查集

题目描述

题目传送门

分析

为了保证字典序最小，我们需要倒序枚举
对于 \(K=1\) 的情况，如果我们当前扫到的值与之前的某个值之和恰好是平方数，那么就要再开一个组
对于 \(k=2\) 的情况，我们把不合法的权值两两连边
那么出现冲突的情况就是整张图不是二分图，也就是说出现了奇环
这个我们用带权并查集随便判一下就可以了
因为我们要对于权值开并查集，所以要注意几种不合法情况的特判
1、当前的值出现的次数大于等于 \(3\)，并且当前值乘 \(2\) 后为平方数
2、当前值与之前某个值冲突，并且这两个值中有一个出现次数大于等于 \(2\)

代码

#include<cstdio>
#include<cmath>
#include<iostream>
#include<algorithm>
#include<map>
#include<vector>
#define rg register
inline int read(){
	rg int x=0,fh=1;
	rg char ch=getchar();
	while(ch<'0' || ch>'9'){
		if(ch=='-') fh=-1;
		ch=getchar();
	}
	while(ch>='0' && ch<='9'){
		x=(x<<1)+(x<<3)+(ch^48);
		ch=getchar();
	}
	return x*fh;
}
const int maxn=4e6+5,maxk=513;
int n,k,a[maxn],sta[maxn],tp,fa[maxn],dis[maxn],mmax;
int mp[maxn];
int zhao(int xx){
	if(xx==fa[xx]) return xx;
	rg int now=fa[xx];
	fa[xx]=zhao(fa[xx]);
	dis[xx]+=dis[now];
	return fa[xx];
}
void bing(int xx,int yy){
	rg int aa=zhao(xx),bb=zhao(yy);
	fa[bb]=aa;
	dis[bb]=dis[xx]-dis[yy]-1;
}
int main(){
	n=read(),k=read();
	if(k==1){
		for(rg int i=1;i<=n;i++){
			a[i]=read();
		}
		rg int now;
		rg bool jud=0;
		for(rg int i=n;i>=1;i--){
			now=std::sqrt(a[i]);
			for(int j=now;j<=maxk;j++){
				if(j*j<a[i]) continue;
				if(mp[j*j-a[i]]){
					jud=1;
					break;
				}
			}
			if(jud){
				sta[++tp]=i;
				if(tp==1){
					for(rg int j=i+1;j<=n;j++){
						mp[a[j]]=0;
					}
				} else {
					for(rg int j=i+1;j<=sta[tp-1];j++){
						mp[a[j]]=0;
					}
				}
				jud=0;
			}
			mp[a[i]]++;
		}
		printf("%d\n",tp+1);
		for(rg int i=tp;i>=1;i--){
			printf("%d ",sta[i]);
		}
		printf("\n");
	} else {
		for(rg int i=1;i<=n;i++){
			a[i]=read();
			mmax=std::max(mmax,a[i]);
		}
		for(rg int i=1;i<=mmax;i++){
			fa[i]=i;
		}
		rg int now;
		rg bool jud=0;
		for(rg int i=n;i>=1;i--){
			now=std::sqrt(a[i]);
			if(mp[a[i]]>=2 && a[i]+a[i]==(int)sqrt(a[i]+a[i])*(int)sqrt(a[i]+a[i])){
				jud=1;
			}
			for(rg int j=now;j<=maxk;j++){
				if(jud) break;
				if(j*j==a[i]+a[i] || j*j<a[i] || j*j-a[i]>mmax) continue;
				if(mp[j*j-a[i]] && mp[a[i]] && a[i]+a[i]==(int)sqrt(a[i]+a[i])*(int)sqrt(a[i]+a[i])){
					jud=1;
					break;
				}
				if(mp[j*j-a[i]]>=2 && j*j-a[i]+j*j-a[i]==(int)sqrt(j*j-a[i]+j*j-a[i])*(int)sqrt(j*j-a[i]+j*j-a[i])){
					jud=1;
					break;
				}
				if(mp[j*j-a[i]]){
					if(zhao(a[i])==zhao(j*j-a[i]) && (dis[a[i]]+dis[j*j-a[i]])%2==0){
						jud=1;
						break;
					} else {
						bing(a[i],j*j-a[i]);
					}
				}
			}
			if(jud){
				sta[++tp]=i;
				if(tp==1){
					for(rg int j=i+1;j<=n;j++){
						mp[a[j]]=0;
						fa[a[j]]=a[j];
						dis[a[j]]=0;
					}
				} else {
					for(rg int j=i+1;j<=sta[tp-1];j++){
						mp[a[j]]=0;
						fa[a[j]]=a[j];
						dis[a[j]]=0;
					}
				}
				jud=0;
			}
			mp[a[i]]++;
		}
		printf("%d\n",tp+1);
		for(rg int i=tp;i>=1;i--){
			printf("%d ",sta[i]);
		}
		printf("\n");
	}
	return 0;
}