P3940 分组

P3940 分组

https://www.luogu.org/problemnew/show/P3940

官方题解http://pan.baidu.com/s/1eSAMuXk

分析：

　　并查集。

　　首先根据K=1和K=2分成两个问题来做。

　　K=1：问题为分成最小数量的区间，使得每个区间满足：任意两个数的和都不是完全平方数（1,3位于一个集合是不可以的），输出字典序最小的方案。

　　一个性质：一个集合越长越好（ppt里有证明）。那么从后往前扫，每到一个判断是否可以加进去，不可以的时候，记录答案。如果直接判断是$n^2$，可以枚举$x^2$来判断，复杂度$O(n\sqrt n)$。

　　K=2：问题为分成最小数量的区间，使得每个区间满足：区间可以任意划分为至多两个集合，这两个集合内部任意两个数的和都不是完全平方数。（1,3,5是可以在分成一段的，因为可以划分为{1,3},{5}）

　　上面的性质还是适用的，所以同样是从后往前扫，每扫到一个判断是否可以加入。这个地方可以记录一个并查集，维护“敌人”集合。$x+n$表示x的“敌人”（不可以一起存在的）。那么如果两个数a+b属于同一个集合，说明a和b用一个共同的敌人，那么a和b就必须在同一个集合里。加入一个x，判断是否可以$k^2-x$共存在这个区间里。注意到一个问题：如果一个数出现了多次，如果它是$2x=k^2$的形式，需要特判（见代码），其他的都可以放进一个集合里。

代码：

 #include<cstdio>
 #include<algorithm>
 #include<cstring>
 #include<cmath>
 #include<iostream>
 #include<cctype>
 #include<set>
 #include<vector>
 #include<queue>
 #include<map>
 #define fi(s) freopen(s,"r",stdin);
 #define fo(s) freopen(s,"w",stdout);
 using namespace std;
 typedef long long LL;

 inline int read() {
     int x=,f=;char ch=getchar();for(;!isdigit(ch);ch=getchar())if(ch=='-')f=-;
     for(;isdigit(ch);ch=getchar())x=x*+ch-'';return x*f;
 }

 const int N = ;

 int a[N], fa[N << ];
 int n, Mx;
 bool vis[N], Ban[N], issqr[N << ];
 vector<int> ans;

 void solve1() {
     for (int j = n, i = n; i>=; ) {
         for (bool flag = ; j >= ; -- j) {
             for (int k=, tmp; ; ++ k) {
                 tmp = k * k - a[j];
                 if (tmp <= ) continue; if (tmp > Mx) break;
                 if (vis[tmp]) { flag = ; break; }
             }
             if (!flag) break;
             vis[a[j]] = true;
         }
         if (!j) break;
         ans.push_back(j);
         for (; i > j; -- i) vis[a[i]] = false;
     }
 }

 int find(int x) {
     return x == fa[x] ? x : fa[x] = find(fa[x]);
 }

 bool check(int u,int v) {
     int u1 = find(u), u2 = find(u + Mx);
     int v1 = find(v), v2 = find(v + Mx);
     if (u1 == v1) return ;
     fa[u1] = v2; fa[v1] = u2;
     return ;
 }

 void solve2() {
     for (int i = ; i <= (Mx << ); ++i) fa[i] = i;
     for (int i = ; i * i <= (Mx << ); ++i) issqr[i * i] = ;
     for (int i = n, j = n; i; ) {
         for (bool flag = ; j; --j) {
             if (!vis[a[j]]) { // 未出现过
                 for (int k = , tmp; ; ++k) {
                     tmp = k * k - a[j];
                     if (tmp <= ) continue; if (tmp > Mx) break;
                     if (vis[tmp] && (Ban[a[j]] || Ban[k * k - a[j]] || check(k * k - a[j], a[j]))) {
                         flag = ; fa[a[j]] = a[j], fa[a[j] + Mx] = a[j] + Mx; break;
                     }
                 }
                 if (!flag) break;
                 vis[a[j]] = true;
             }
             // 出现过：只有2*a=x^2（2,8...）的情况要特判，其他的都可以分到一集合里。
             // 这些数最多出现两个，而且没有第三个敌人
             else if (issqr[a[j] + a[j]]) {
                 if (Ban[a[j]]) break;
                 for (int k = , tmp; ; ++k) { // 判断是否有第三个敌人
                     tmp = k * k - a[j];
                     if (tmp < ) continue; if (tmp > Mx) break;
                     if (vis[tmp] && k * k != a[j] * ) { flag = ; break; }
                 }
                 if (!flag) break;
                 Ban[a[j]] = true; // a[j]以前一个，现在一个 a[j]以后不能再有了。
             }
         }
         if (!j) break;
         ans.push_back(j);
         for (; i > j; --i) fa[a[i]] = a[i], fa[a[i] + Mx] = a[i] + Mx, vis[a[i]] = Ban[a[i]] = ;
     }
 }

 int main() {
     n = read();int K = read();
     for (int i=; i<=n; ++i) a[i] = read(), Mx = max(Mx, a[i]);
     if (K == ) solve1();
     else solve2();
     printf("%d\n",ans.size() + );
     for (int i=ans.size()-; i>=; --i) printf("%d ",ans[i]); putchar('\n');
     return ;
 }