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

代码：

#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=0,f=1;char ch=getchar();for(;!isdigit(ch);ch=getchar())if(ch=='-')f=-1;
    for(;isdigit(ch);ch=getchar())x=x*10+ch-'0';return x*f;
}

const int N = 131073;

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

void solve1() {
    for (int j = n, i = n; i>=1; ) {
        for (bool flag = 1; j >= 1; -- j) {
            for (int k=1, tmp; ; ++ k) {
                tmp = k * k - a[j];
                if (tmp <= 0) continue; if (tmp > Mx) break;
                if (vis[tmp]) { flag = 0; 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 1;
    fa[u1] = v2; fa[v1] = u2;
    return 0;
}

void solve2() {
    for (int i = 1; i <= (Mx << 1); ++i) fa[i] = i;
    for (int i = 1; i * i <= (Mx << 1); ++i) issqr[i * i] = 1;
    for (int i = n, j = n; i; ) {
        for (bool flag = 1; j; --j) {
            if (!vis[a[j]]) { // 未出现过 
                for (int k = 1, tmp; ; ++k) {
                    tmp = k * k - a[j];
                    if (tmp <= 0) continue; if (tmp > Mx) break;
                    if (vis[tmp] && (Ban[a[j]] || Ban[k * k - a[j]] || check(k * k - a[j], a[j]))) {
                        flag = 0; 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 = 1, tmp; ; ++k) { // 判断是否有第三个敌人 
                    tmp = k * k - a[j];
                    if (tmp < 0) continue; if (tmp > Mx) break;
                    if (vis[tmp] && k * k != a[j] * 2) { flag = 0; 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]] = 0;
    }
}

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