luogu3940 分组 题解

1
// Problem: P3940 分组
2
// Contest: Luogu
3
// URL: https://www.luogu.com.cn/problem/P3940
4
// Author: yozora0908
5
// Memory Limit: 250 MB
6
// Time Limit: 1000 ms
7
//
8
// Let's Daze
9
//
10
// Powered by CP Editor (https://cpeditor.org)
11

12
#include<bits/stdc++.h>
13
using namespace std;
14
#define ll long long
15
#define uint unsigned long long
16
#define PII pair<int,int>
17
#define MP make_pair
18
#define fi first
19
#define se second
20
#define pb push_back
21
#define eb emplace_back
22
#define SET(a,b) memset(a,b,sizeof(a))
23
#define CPY(a,b) memcpy(a,b,sizeof(b))
24
#define rep(i,j,k) for(int i=(j);i<=(k);++i)
25
#define per(i,j,k) for(int i=(j);i>=(k);--i)
26
int read() {
27
    int a=0, f=1; char c=getchar();
28
    while(!isdigit(c)) {
29
        if(c=='-') f=-1;
30
        c=getchar();
31
    }
32
    while(isdigit(c)) a=a*10+c-'0', c=getchar();
33
    return a*f;
34
}
35
const int N=131075, lim=131072;
36
int n, m, K, a[N], ans[N], sqr[2*N];
37
bool v[N], v2[N];
38
namespace sub1 {
39
    void solve() {
40
        for(int i=n,j=n;i;) {
41
            for(;j;--j) {
42
                for(int k=1;k*k-a[j]<N;++k) if(a[j]<=k*k) {
43
                    if(v[k*k-a[j]]) goto out;
44
                }
45
                v[a[j]]=1;
46
            }
47
            out:;
48
            if(!j) break;
49
            ans[++m]=j;
50
            for(;i>j;--i) v[a[i]]=0;
51
        }
52

53
    }
54
};
55
struct DSU {
56
    int f[2*N];
57
    void init() { for(int i=1;i<=2*lim;++i) f[i]=i; }
58
    int get(int x) { return x==f[x]? x:f[x]=get(f[x]); }
59
    void merge(int x,int y) {
60
        x=get(x), y=get(y);
61
        if(x!=y) f[x]=y;
62
    }
63

64
} dsu;
65
bool check(int x,int y) {
66
    int x1=dsu.get(x), x2=dsu.get(x+lim);
67
    int y1=dsu.get(y), y2=dsu.get(y+lim);
68
    if(x1==y1) return 1;
69
    if(x2==y2) return 1;
70
    dsu.merge(x1,y2);
71
    dsu.merge(x2,y1);
72
    return 0;
73
}
74
namespace sub2 {
75
    void solve() {
76
        dsu.init();
77
        for(int i=n,j=n;i;) {
78
            for(;j;--j) {
79
                if(!v[a[j]]) {
80
                    for(int k=1;k*k-a[j]<N;++k) if(a[j]<=k*k) {
81
                        if(v[k*k-a[j]]) {
82
                            if(check(a[j],k*k-a[j])||v2[k*k-a[j]]) goto out;
83
                        }
84
                    }
85
                    v[a[j]]=1;
86
                } else {
87
                    if(sqr[2*a[j]]) {
88
                        if(v2[a[j]]) goto out;
89
                        v2[a[j]]=1;
90
                        for(int k=1;k*k-a[j]<N;++k) if(a[j]<=k*k) {
91
                            if(v[k*k-a[j]]&&k*k!=2*a[j])  goto out;
92
                        }
93

94
                    }
95
                }
96
            }
97
            out:;
98
            if(!j) break;
99
            ans[++m]=j;
100
            for(;i>j;--i) v[a[i]]=v2[a[i]]=0, dsu.f[a[i]]=a[i], dsu.f[a[i]+lim]=a[i]+lim;
101

102
        }
103
    }
104
};
105
signed main() {
106
    n=read(), K=read();
107
    rep(i,1,n) a[i]=read();
108
    for(int i=1;i<=512;++i) sqr[i*i]=1;
109
    if(K==1) sub1::solve();
110
    else sub2::solve();
111
    printf("%d\n",m+1);
112
    per(i,m,1) printf("%d ",ans[i]);
113
    return 0;
114
}