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(nsqrt 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$的形式，需要特判（见代码），其他的都可以放进一个集合里。

代码：

  1 #include<cstdio>
  2 #include<algorithm>
  3 #include<cstring>
  4 #include<cmath>
  5 #include<iostream>
  6 #include<cctype>
  7 #include<set>
  8 #include<vector>
  9 #include<queue>
 10 #include<map>
 11 #define fi(s) freopen(s,"r",stdin);
 12 #define fo(s) freopen(s,"w",stdout);
 13 using namespace std;
 14 typedef long long LL;
 15 
 16 inline int read() {
 17     int x=0,f=1;char ch=getchar();for(;!isdigit(ch);ch=getchar())if(ch=='-')f=-1;
 18     for(;isdigit(ch);ch=getchar())x=x*10+ch-'0';return x*f;
 19 }
 20 
 21 const int N = 131073;
 22 
 23 int a[N], fa[N << 1];
 24 int n, Mx;
 25 bool vis[N], Ban[N], issqr[N << 1];
 26 vector<int> ans;
 27 
 28 void solve1() {
 29     for (int j = n, i = n; i>=1; ) {
 30         for (bool flag = 1; j >= 1; -- j) {
 31             for (int k=1, tmp; ; ++ k) {
 32                 tmp = k * k - a[j];
 33                 if (tmp <= 0) continue; if (tmp > Mx) break;
 34                 if (vis[tmp]) { flag = 0; break; }
 35             }
 36             if (!flag) break;
 37             vis[a[j]] = true;
 38         }
 39         if (!j) break;
 40         ans.push_back(j);
 41         for (; i > j; -- i) vis[a[i]] = false;
 42     }
 43 }
 44 
 45 int find(int x) {
 46     return x == fa[x] ? x : fa[x] = find(fa[x]);
 47 }
 48 
 49 bool check(int u,int v) {
 50     int u1 = find(u), u2 = find(u + Mx);
 51     int v1 = find(v), v2 = find(v + Mx);
 52     if (u1 == v1) return 1;
 53     fa[u1] = v2; fa[v1] = u2;
 54     return 0;
 55 }
 56 
 57 void solve2() {
 58     for (int i = 1; i <= (Mx << 1); ++i) fa[i] = i;
 59     for (int i = 1; i * i <= (Mx << 1); ++i) issqr[i * i] = 1;
 60     for (int i = n, j = n; i; ) {
 61         for (bool flag = 1; j; --j) {
 62             if (!vis[a[j]]) { // 未出现过 
 63                 for (int k = 1, tmp; ; ++k) {
 64                     tmp = k * k - a[j];
 65                     if (tmp <= 0) continue; if (tmp > Mx) break;
 66                     if (vis[tmp] && (Ban[a[j]] || Ban[k * k - a[j]] || check(k * k - a[j], a[j]))) {
 67                         flag = 0; fa[a[j]] = a[j], fa[a[j] + Mx] = a[j] + Mx; break; 
 68                     }
 69                 }
 70                 if (!flag) break;
 71                 vis[a[j]] = true;
 72             }
 73             // 出现过：只有2*a=x^2（2,8...）的情况要特判，其他的都可以分到一集合里。
 74             // 这些数最多出现两个，而且没有第三个敌人 
 75             else if (issqr[a[j] + a[j]]) {
 76                 if (Ban[a[j]]) break;
 77                 for (int k = 1, tmp; ; ++k) { // 判断是否有第三个敌人 
 78                     tmp = k * k - a[j];
 79                     if (tmp < 0) continue; if (tmp > Mx) break;
 80                     if (vis[tmp] && k * k != a[j] * 2) { flag = 0; break; }
 81                 }
 82                 if (!flag) break;
 83                 Ban[a[j]] = true; // a[j]以前一个，现在一个 a[j]以后不能再有了。 
 84             }
 85         }
 86         if (!j) break;
 87         ans.push_back(j);
 88         for (; i > j; --i) fa[a[i]] = a[i], fa[a[i] + Mx] = a[i] + Mx, vis[a[i]] = Ban[a[i]] = 0;
 89     }
 90 }
 91 
 92 int main() {
 93     n = read();int K = read();
 94     for (int i=1; i<=n; ++i) a[i] = read(), Mx = max(Mx, a[i]);
 95     if (K == 1) solve1();
 96     else solve2();
 97     printf("%d
",ans.size() + 1);
 98     for (int i=ans.size()-1; i>=0; --i) printf("%d ",ans[i]); putchar('
');
 99     return 0;
100 }