題解 P3940 分組
阿新 • • 發佈:2021-06-20
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; }