题解 P3940 分组

Solution

对于字典序最小 , 从后往前扫即可满足 .
对于查询冲突 , 通过枚举完全平方数解决 .
\(K=1\) 时 , 出现冲突时新分一组即可 .
\(K=2\) 时 , 通过拓展域并查集来判断 .
我们把一个数拆成两个点 , 分别为黑点和白点 , 对于冲突的点 , 将它们的黑白点两两相连 , 注意特判 \(2a_ i=x^2\) 的情况
时间复杂度 \(O(n\sqrt a)\)

Code

#include<iostream>
#include<cstdio>
#include<cmath>
#include<unordered_map>
using namespace std;
int read()
{
	int ret=0;char c=getchar();
	while(c>'9'||c<'0')c=getchar();
	while(c>='0'&&c<='9')ret=(ret<<3)+(ret<<1)+(c^48),c=getchar();
	return ret;
}
const int maxn=131073; 
int n,k;
int a[maxn],maxa;
int pown[550],exi[maxn];
bool s[maxn];int use=1;
bool judge(int x){int rt=sqrt(x);return (rt*rt)==x;}
struct dsu
{
    int fa[maxn<<1];
    void prework(){for(int i=1;i<2*maxn;i++)fa[i]=i;}
    int get(int x)
    {
        if(x==fa[x])return fa[x];
        return fa[x]=get(fa[x]);
    }
    void merge(int x,int y){fa[get(x)]=get(y);}
}d;
void end(int x)
{
    use++;s[x]=1;
    int now=x+1;
    while(!s[now]&&now<=n)
    {
        d.fa[a[now]]=a[now];
        d.fa[a[now]+maxn]=a[now]+maxn;
        exi[a[now++]]=0;
    }
    d.fa[a[now]]=a[now];
    d.fa[a[now]+maxn]=a[now]+maxn;
    exi[a[now]]=0;
    d.fa[a[x]]=a[x];
    d.fa[a[x]+maxn]=a[x]+maxn;
    exi[a[x]]=1;
}
unordered_map<int,int>exi1;
int main()
{
    n=read();k=read();
    for(int i=1;i<=n;i++)maxa=max(maxa,a[i]=read());
    d.prework();
    maxa=ceil(sqrt(2*maxa));
    for(int i=1;i<=maxa;i++)pown[i]=i*i;
    if(k==1)
    {
        for(int i=n;i>=1;i--)
        {
            bool flag=0;
            for(int j=1;j<=maxa;j++)
            {
                if(exi1.find(pown[j]-a[i])!=exi1.end()){flag=1;break;}
            }
            if(flag){use++;s[i]=1;exi1.clear();}        
            exi1[a[i]]++;
        }
    }
    else if(k==2)
    {
        for(int i=n;i>=1;i--)
        {
            if(judge(2*a[i])&&exi[a[i]])
            {
                if(exi[a[i]]==2)end(i);
                else if(exi[a[i]]==1)
                {                
                    bool flag=0;
                    for(int j=maxa;j>=1;j--)
                    {
                        if(pown[j]-a[i]==a[i])continue;
                        if(pown[j]-a[i]<=0)break;
                        if(pown[j]-a[i]>=maxn)continue;
                        if(exi[pown[j]-a[i]]){flag=1;break;}
                    }
                    if(!flag)exi[a[i]]++;
                    else end(i);
                }
               continue;
            }
            if(exi[a[i]]){exi[a[i]]++;continue;}
            for(int j=maxa;j>=1;j--)
            {
                if(pown[j]-a[i]<=0)break;
                if(pown[j]-a[i]>=maxn)continue;
                if(exi[pown[j]-a[i]])
                {
                    d.merge(a[i],pown[j]-a[i]+maxn);
                    d.merge(a[i]+maxn,pown[j]-a[i]);
                    if(d.get(a[i])==d.get(a[i]+maxn))break;
                }
                if(exi[pown[j]-a[i]]==2&&judge(2*(pown[j]-a[i]))){d.merge(a[i],a[i]+maxn);break;}
            }
            if(d.get(a[i])==d.get(a[i]+maxn))end(i);
            else exi[a[i]]++;
        }
    }
    printf("%d\n",use);
    for(int i=1;i<n;i++){if(s[i])printf("%d ",i);}
    printf("\n");
    return 0;
}