[洛谷P3940]:分组（贪心+并查集）

题目传送门

---

题目描述

小$C$在了解了她所需要的信息之后，让兔子们调整到了恰当的位置。小$C$准备给兔子们分成若干个小组来喂恰当的胡萝卜给兔子们吃。
此时，$n$只兔子按一定顺序排成一排，第$i$只兔子的颜色是$a_i$。由于顺序已经是被调整好了的，所以每个小组都应当是序列上连续的一段。
在分组前，小$C$发现了一个规律：有些兔子会两两发生矛盾。并且，两只兔子会发生矛盾，当且仅当代表他们的颜色的数值之和为一个正整数的平方。比如，$1$色兔子和$2$色兔子不会发生矛盾，因为$3$不是任何一个正整数的平方；而$1$色兔子却会和$3$色兔子发生矛盾，因为$4=2^2$。
小$C$认为，只要一个小组内的矛盾不要过大就行。因此，小$C$定义了一个小组的矛盾值$k$，表示在这个小组里，至少需要将这个组再一次分成$k$个小团体；每个小团体并不需要是序列上连续的一段，但是需要使得每个小团体内任意两只兔子之间都不会发生矛盾。
小$C$要求，矛盾值最大的小组的矛盾值$k$不超过$K$就可以了。当然，这样的分组方法可能会有很多个；为了使得分组变得更加和谐，小$C$想知道，在保证分组数量最少的情况下，字典序最小的方案是什么。你能帮帮她吗？
字典序最小的方案是指，按顺序排列分组的间隔位置，即所有存在兔子$i$和$i+1$在不同组的位置。

---

输入格式

输入第1行两个正整数$n,K$。
输入第2行$n$个正整数，第$i$个数表示第$i$只兔子的颜色$a$。

---

输出格式

输出第$1$行一个正整数$m$，为你至少需要将兔子分为多少个小组。
输出第$2$行$m-1$个从小到大的排列的正整数，第$i$个数$s_i$表示$s_i$和$s_{i+1}$在你的方案里被分到了两个小组。如果$m=1$，那么请输出一个空行。

---

样例

样例输入1：

5 2
1 3 15 10 6

样例输出1：

2
1

样例输入2：

319 2
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319

样例输出2：

6
7 31 71 127 199

---

数据范围与提示

样例$1$解释：

如果将五只兔子全部分到同一个小组的话，那么$(1,3)(3,6)(6,10)(10,15)(1,15)$均不能分到同一个小团体；因为最多分成两个小团体，所以为了满足前4对限制，只能分为$\{\{1,6,15\},\{3,10\}\}$，但此时不满足$(1,15)$，所以不存在一种组数为$1$的方案满足全部限制。
如果将五只兔子分为两个小组的话，一种字典序最小的可行的分组方案是$\{1\},\{3,15,10,6\}$，此时第二组内的小团体数量不超过$2$的一种分法是$\{\{3,10\},\{15,6\}\}$。

数据范围：

任务会给出部分测试数据的特点。如果你在解决题目中遇到了困难，可以尝试只解决一部分测试数据。
每个测试点的数据规模及特点如下表：

测试点       $K$           $n$                 $a_i,c_j$               特殊性质$1$           特殊性质$2$
    $1$             $1$           $2$                 $\leqslant 131072$              $\surd$                          $\times$
    $2$             $1$           $4$                 $\leqslant 131072$              $\times$                          $\times$
    $3$             $1$          $16$               $\leqslant 131072$               $\times$                          $\times$
    $4$             $1$          $256$             $\leqslant  131072$              $\surd$                          $\times$
    $5$             $1$          $1024$          $\leqslant 131072$              $\surd$                          $\times$
    $6$             $1$          $1024$          $\leqslant 131072$              $\times$                          $\times$
    $7$             $1$          $131072$     $\leqslant 2$                          $\surd$                          $\times$
    $8$             $1$          $131072$     $\leqslant 131072$              $\surd$                          $\times$
    $9$             $1$          $131072$     $\leqslant 131072$              $\times$                          $\surd$
    $10$           $1$         $131072$      $\leqslant 131072$              $\times$                          $\times$
    $11$           $2$         $2$                  $\leqslant 131072$              $\surd$                          $\times$
    $12$           $2$         $4$                  $\leqslant 131072$              $\times$                          $\times$
    $13$           $2$         $8$                  $\leqslant 131072$              $\times$                          $\times$
    $14$           $2$         $16$               $\leqslant 131072$               $\times$                          $\times$
    $15$           $2$         $256$             $\leqslant 131072$              $\surd$                          $\times$
    $16$           $2$         $256$             $\leqslant 131072$              $\times$                          $\surd$
    $17$           $2$         $1024$           $\leqslant 2$                          $\surd$                          $\times$
    $18$           $2$         $1024$           $\leqslant 131072$              $\times$                          $\times$
    $19$           $2$         $4096$           $\leqslant 2$                          $\times$                           $\times$
    $20$           $2$         $4096$           $\leqslant 131072$              $\times$                          $\times$
    $21$           $2$         $131072$      $\leqslant 2$                           $\times$                          $\times$
    $22$           $2$         $131072$      $\leqslant 131072$              $\times$                          $\surd$
    $23$           $2$         $131072$      $\leqslant 131072$              $\times$                          $\surd$
    $24$           $2$         $13102$        $\leqslant 131072$              $\times$                          $\times$
    $25$           $2$         $131072$      $\leqslant 131072$              $\times$                          $\times$

---

题解

对于$K=1$：

　测试点$1$：$n=2$

　　直接判一下两数值和是否为完全平方数。

　　是就$puts("2$ \ $n1");$，不是就$puts("1$ \ $n");$。

　　时间复杂度：$\Theta(1)$。

　　期望得分：$4$分。

　　实际得分：$4$分。

　　总得分：$4$分。

　测试点$2$：$n=4$

　　手动讨论所有情况即可。

　　时间复杂度：$\Theta(1)$。

　　期望得分：$4$分。

　　实际得分：$4$分。

　　总得分：$8$分。

　测试点$3$：$n=16$

　　暴力搜索，可以得到所有的分组情况，一个一个验证就好了。

　　时间复杂度：$\Theta(2^n\times n^2)$。

　　期望得分：$12$分。

　　实际得分：$12$分。

　　总得分：$12$分。

　测试点$4$：$n=256$

　　考虑$DP$。

　　预处理$f[i][j]$表示从$i$到$j$可以分到一组。

　　定义$dp[i][j]$表示，此时已经分到了第$i$个位置，已经分了$k$组，方案可行$or$不可行（$0$/$1$表示。

　　则$DP$式子为：$dp[i][k]|=dp[j][k-1]\&\&f[j+1][i]$。

　　在开一个数组记录一下上次转移的位置，为保证最优性，已经转移过不再更新，判断一下就好了。

　　时间复杂度：$\Theta(n^3)$。

　　期望得分：$16$分。

　　实际得分：$24$分。

　　总得分：$24$分。

　测试点$5,6$：$n=1024$

　　发现上面的解法中$k$那一维可以去掉，即：$if(f[j+1][i])dp[i]=\min(dp[i],dp[j]+1)$。

　　时间复杂度：$\Theta(n^2)$。

　　期望得分：$24$分。

　　实际得分：$24$分。

　测试点$7\thicksim 10$：$n=131072$

　　根据我们刚刚的结论，细心的你一定会发现：

　　其实这个$DP$并不是一个$DP$，本质上是找最长的合法分段，然后把问题转化为子问题。

　　问题可以转化为：从前往后，选择一段最长的合法区间并分割，重新进行知道完成为止。

　　为什么要从后往前呢？因为我们要保证在最优的情况下断点的字典序最小，所以我们需要让断点尽可能的靠前，然而从前往后搜我们并不知道什么时候该断才是最优的，所以我们就可以从后往前搜，尽可能的靠后断，就能保证答案最优了。

　　我们可以在每插入一个点时逐一判断是否不合法。

　　时间复杂度：$\Theta(n^2)$。

　　期望得分：$24$分。

　　实际得分：$32$分。

　　总得分：$32$分。

　　显然上面的这种方法会让你$T$到飞起，所以考虑如何优化。

　　悄悄的观察数据范围会惊喜的发现$a_i\leqslant 131072$，那么就意味着$a_i+a_j=262144={512}^2$，所以我们在插入一个数的时候可以枚举$x$，看$x^2-a_j$有没有出现过即可。

　　时间复杂度：$\Theta(n\sqrt{n})$。

　　期望得分：$40$分。

　　实际得分：$40$分。

　　总得分：$40$分。

这样，我们对$K=1$的情况就拿满分了。

对于$K==2$：

　测试点$11$：$n=2$

　　无论如何，一组也够了（出题人好良心……）。

　　直接$puts("1$\ $n");$就好啦。

　　时间复杂度：$\Theta(1)$。

　　期望得分：$4$分。

　　实际得分：$4$分。

　　总得分：$44$分。

　测试点$12\thicksim 14$：

　　枚举分组，判定是否合法。

　　如何判定合法？

　　可以再次进行暴力分组。

　　时间复杂度：$\Theta(3^n)$。

　　其实没有必要，把会发生矛盾的兔子连边，发现这就是个二分图判定，黑白染色即可。

　　时间复杂度：$\Theta(2^n\times n)$。

　　期望得分：$16$分。

　　实际得分：$16$分。

　　总得分：$56$分。

　测试点$15\thicksim 18$：

　　根据之前的结论，发现我们就是找判定最长序列。

　　每次判定时建一个图进行判定，直到非法为止。

　　时间复杂度：$\Theta(n^3)$。

　　期望得分：$24$分。

　　实际得分：$32$分。

　　总得分：$72$分。

　测试点$19\thicksim 25$：

　　发现我们上面的用二维数组记录敌对关系是一种限制，所以我们可以考虑使用带权并查集来存储关系。

　　可以仿照$K=1$的思想进行统计答案。

　　但是需要注意一个问题：

　　　　如果某个数出现了多次，如果这个数的二倍是一个完全平方数，那么就不能丢到同一组；否则不管有多少个也不用担心（注意还不能有第三个数与它们敌对）。

　　时间复杂度：$\Theta(n\times \sqrt{\max(a_i)})$。

　　期望得分：$60$分。

　　实际得分：$60$分。

　　总得分：$100$分。

到这里，我们就成功的解决了这道题～～～

---

代码时刻

测试点$1$：

#include<bits/stdc++.h>
using namespace std;
int a[131073];
int main()
{
	int n,K;
	scanf("%d%d",&n,&K);
	for(int i=1;i<=n;i++)
		scanf("%d",&a[i]);
	if(K==1)
	{
		if(n==2)
			if((int)sqrt(a[1]+a[2])*(int)sqrt(a[1]+a[2])==a[1]+a[2])puts("2\n1");
			else puts("1\n");
	}		
	return 0;
}

测试点$2$：

#include<bits/stdc++.h>
using namespace std;
int a[131073];
bool asksqr(int x,int y){return (int)sqrt(x+y)*(int)sqrt(x+y)==x+y?0:1;}
void judge0()
{
	if(asksqr(a[1],a[2])&&asksqr(a[1],a[3])&&asksqr(a[1],a[4])&&asksqr(a[2],a[3])&&asksqr(a[2],a[4])&&asksqr(a[3],a[4])){puts("1\n");exit(0);}
}
void judge1()
{
	if(asksqr(a[2],a[3])&&asksqr(a[2],a[4])&&asksqr(a[3],a[4])){puts("2\n1");exit(0);}
	if(asksqr(a[1],a[2])&&asksqr(a[3],a[4])){puts("2\n2");exit(0);}
	if(asksqr(a[1],a[2])&&asksqr(a[1],a[3])&&asksqr(a[2],a[3])){puts("2\n3");exit(0);}
}
void judge2()
{
	if(asksqr(a[3],a[4])){puts("3\n1\n2");exit(0);}
	if(asksqr(a[2],a[3])){puts("3\n1\n3");exit(0);}
	if(asksqr(a[1],a[2])){puts("3\n3\n3");exit(0);}
}
void judge3(){puts("4\n1\n2\n3");exit(0);}
int main()
{
	int n,K;
	scanf("%d%d",&n,&K);
	for(int i=1;i<=n;i++)
		scanf("%d",&a[i]);
	if(K==1)
	{
		if(n==4)
		{
			judge0();
			judge1();
			judge2();
			judge3();
		}
	}		
	return 0;
}

测试点$3$：

#include<bits/stdc++.h>
using namespace std;
int n,K;
int a[131073];
int ans[131072],cnt;
bool asksqr(int x,int y){return (int)sqrt(x+y)*(int)sqrt(x+y)==x+y?1:0;}
bool judge()
{
	for(int i=1;i<=cnt;i++)
		for(int j=ans[i-1]+1;j<ans[i];j++)
			for(int k=j+1;k<=ans[i];k++)
				if(asksqr(a[j],a[k]))return 0;
	for(int i=ans[cnt]+1;i<n;i++)
		for(int j=i+1;j<=n;j++)
			if(asksqr(a[i],a[j]))return 0;
	return 1;
}
void print()
{
	printf("%d\n",cnt+1);
	for(int i=1;i<=cnt;i++)
		printf("%d ",ans[i]);
}
void dfs(int x,int l,int w)
{
	if(x==w){if(judge()){print();exit(0);}return;}
	for(int i=l;i<n;i++)
	{
		ans[++cnt]=i;
		dfs(x+1,i+1,w);
		cnt--;
	}
}
int main()
{
	scanf("%d%d",&n,&K);
	for(int i=1;i<=n;i++)
		scanf("%d",&a[i]);
	if(K==1)
	{
		for(int i=0;i<n;i++)dfs(0,1,i);
	}
	return 0;
}

测试点$4$：

#include<bits/stdc++.h>
using namespace std;
int n,K;
int a[131073];
bool f[1050][1050],dp[1050][1050];
int flag[1050][1050];
int ans[1050];
bool asksqr(int x,int y){return (int)sqrt(x+y)*(int)sqrt(x+y)==(x+y)?0:1;}
int main()
{
	scanf("%d%d",&n,&K);
	for(int i=1;i<=n;i++)
	{
		scanf("%d",&a[i]);
		dp[i][i]=1;
	}
	dp[0][0]=1;
	for(int i=1;i<=n;i++)
	{
		for(int j=i;j<=n;j++)
		{
			for(int k=i;k<j;k++)
				if(!asksqr(a[j],a[k]))goto nxt;
			f[i][j]=f[j][i]=1;
		}
		nxt:;
	}
	for(int i=1;i<=n;i++)
		for(int j=0;j<i;j++)
			for(int k=1;k<=n;k++)
			{
				dp[i][k]|=dp[j][k-1]&f[j+1][i];
				if(dp[j][k-1]&&f[j+1][i]&&!flag[i][k])flag[i][k]=j;
			}
	for(int i=1;i<=n;i++)
		if(dp[n][i])
		{
			printf("%d\n",i);
			int now=n;
			int k=i;
			while(now)
			{
				ans[++ans[0]]=flag[now][k];
				now=flag[now][k--];
			}
			break;
		}
	for(int i=ans[0]-1;i;i--)
		printf("%d ",ans[i]);
	return 0;
}

测试点$5,6$：

#include<bits/stdc++.h>
using namespace std;
int n,K;
int a[131073];
bool f[1050][1050];
int dp[1050];
int flag[1050][1050];
int ans[1050];
bool asksqr(int x,int y){return (int)sqrt(x+y)*(int)sqrt(x+y)==(x+y)?0:1;}
int main()
{
	scanf("%d%d",&n,&K);
	memset(dp,0x3f,sizeof(dp));
	for(int i=1;i<=n;i++)
		scanf("%d",&a[i]);
	dp[0]=1;
	for(int i=1;i<=n;i++)
	{
		for(int j=i;j<=n;j++)
		{
			for(int k=i;k<j;k++)
				if(!asksqr(a[j],a[k]))goto nxt;
			f[i][j]=f[j][i]=1;
		}
		nxt:;
	}
	for(int i=1;i<=n;i++)
		for(int j=0;j<i;j++)
			if(f[j+1][i])
			{
				dp[i]=min(dp[i],dp[j]+1);
				if(!flag[i][dp[j]+1])flag[i][dp[j]+1]=j;
			}
	printf("%d\n",dp[n]-1);
	int now=n;
	int k=dp[n];
	while(now)
	{
		ans[++ans[0]]=flag[now][k];
		now=flag[now][k--];
	}
	for(int i=ans[0]-1;i;i--)
		printf("%d ",ans[i]);
	return 0;
}

测试点$7\thicksim 10$（$32$分）：

#include<bits/stdc++.h>
using namespace std;
int n,K;
int a[150000];
bool flag[300000];
int que[150000],sum;
int main()
{
	scanf("%d%d",&n,&K);
	for(int i=1;i<=n;i++)scanf("%d",&a[i]);
	for(int i=1;i<=512;i++)flag[i*i]=1;
	if(K==1)
	{
		int l=n;
		for(int i=n;i;i--)
			for(int j=l;j>i;j--)
				if(flag[a[i]+a[j]])
				{
					l=i;
					que[++sum]=i;
					break;
				}
		printf("%d\n",sum+1);
		for(int i=sum;i;i--)printf("%d ",que[i]);
	}
	return 0;
}

测试点$7\thicksim 10$（$40$分）：

#include<bits/stdc++.h>
using namespace std;
int n,K;
int a[150000],maxn;
bool vis[150000];
int que[150000],sum;
int main()
{
	scanf("%d%d",&n,&K);
	for(int i=1;i<=n;i++){scanf("%d",&a[i]);maxn=max(maxn,a[i]);}
	if(K==1)
	{
		for(int i=n,j=n;i;)
		{
			for(;j;j--)
			{
				for(int k=1;k*k-a[j]<=maxn;k++)
				{
					if(k*k-a[j]<=0)continue;
					if(vis[k*k-a[j]])goto nxt;
				}
				vis[a[j]]=1;
			}
			nxt:
			if(!j)break;
			que[++sum]=j;
			for(;i>j;i--)vis[a[i]]=0;
		}	
		printf("%d\n",sum+1);
		for(int i=sum;i;i--)printf("%d ",que[i]);
	}
	return 0;
}

测试点$11$：

#include<bits/stdc++.h>
using namespace std;
int n,K;
int a[150000];
int main()
{
	scanf("%d%d",&n,&K);
	for(int i=1;i<=n;i++)
		scanf("%d",&a[i]);
	if(K==2)if(n==2)puts("1\n");
	return 0;
}

测试点$12\thicksim 15$（二分图）：

#include<bits/stdc++.h>
using namespace std;
int n,K;
int a[131073];
int ans[131072],col[131072],que[131072],cnt;
bool Map[30][30],flag[30],vis[30];
bool asksqr(int x,int y){return (int)sqrt(x+y)*(int)sqrt(x+y)==x+y?1:0;}
bool color_solve(int x,int color)//判定二分图
{
	col[x]=color;
	vis[x]=1;
	for(int i=1;i<=n;i++)
	{
		if(i==x||!flag[i]||!Map[x][i])continue;
		if(col[i]==color)return 1;
		if(!col[i]&&color_solve(i,-color))return 1;
	}
	return 0;
}
bool judge()
{
    for(int i=1;i<=cnt;i++)
    {
        que[0]=0;
        memset(vis,0,sizeof(vis));
        memset(flag,0,sizeof(flag));
        memset(col,0,sizeof(col));
        for(int j=ans[i-1]+1;j<=ans[i];j++)
        {que[++que[0]]=j;flag[j]=1;}
        for(int j=1;j<=que[0];j++)
        	if(!vis[que[j]])if(color_solve(que[j],1))return 0;
    }
    que[0]=0;
    memset(vis,0,sizeof(vis));
    memset(flag,0,sizeof(flag));
    memset(col,0,sizeof(col));
    for(int i=ans[cnt]+1;i<=n;i++)
    {que[++que[0]]=i;flag[i]=1;}
    for(int i=1;i<=que[0];i++)
        if(!vis[que[i]])if(color_solve(que[i],1))return 0;
    return 1;
}
void print()
{
    printf("%d\n",cnt+1);
    for(int i=1;i<=cnt;i++)
        printf("%d ",ans[i]);
}
void dfs(int x,int l,int w)
{
    if(x==w){if(judge()){print();exit(0);}return;}
    for(int i=l;i<n;i++)
    {
        ans[++cnt]=i;
        dfs(x+1,i+1,w);
        cnt--;
    }
}
int main()
{
    scanf("%d%d",&n,&K);
    for(int i=1;i<=n;i++)
        scanf("%d",&a[i]);
    for(int i=1;i<n;i++)
    	for(int j=i+1;j<=n;j++)
    		Map[i][j]=Map[j][i]=asksqr(a[i],a[j]);
    if(K==2)
    {
        for(int i=0;i<n;i++)
        	dfs(0,1,i);
    }
    return 0;
}

测试点$15\thicksim 18$：

#include<bits/stdc++.h>
using namespace std;
int n,K;
int a[131073];
int ans[131072],col[131072],que[131072],sum;
bool Map[1050][1050],flag[1050],vis[1050];
bool asksqr(int x,int y){return (int)sqrt(x+y)*(int)sqrt(x+y)==x+y?1:0;}
bool color_solve(int x,int color)
{
	col[x]=color;
	vis[x]=1;
	for(int i=1;i<=n;i++)
	{
		if(i==x||!flag[i]||!Map[x][i])continue;
		if(col[i]==color)return 1;
		if(!col[i]&&color_solve(i,-color))return 1;
	}
	return 0;
}
int main()
{
    scanf("%d%d",&n,&K);
    for(int i=1;i<=n;i++)
        scanf("%d",&a[i]);
    for(int i=1;i<n;i++)
    	for(int j=i+1;j<=n;j++)
    		Map[i][j]=Map[j][i]=asksqr(a[i],a[j]);
    if(K==2)
    {
    	int l=n;
        for(int i=n;i;i--)
        {
		    que[0]=0;
		    memset(vis,0,sizeof(vis));
		    memset(flag,0,sizeof(flag));
		    memset(col,0,sizeof(col));
        	for(int j=l;j>=i;j--)
        	{
        		que[++que[0]]=j;
        		flag[j]=1;
        	}
        	for(int j=1;j<=que[0];j++)
        		if(!vis[que[j]])
        		if(color_solve(que[j],1))
        		{
        			l=i;
        			ans[++sum]=i;
        			break;
        		}
        }
        printf("%d\n",sum+1);
        for(int i=sum;i;i--)printf("%d ",ans[i]);
    }
    return 0;
}

测试点$19\thicksim 25$：

#include<bits/stdc++.h>
using namespace std;
int n,K;
int a[150000],maxn;
bool vis[150000],sit[150000];
int que[150000],sum;
int fa[300000];
int find(int x){return fa[x]<=0?x:fa[x]=find(fa[x]);}
void connect(int x,int y)
{
	x=find(x);
	y=find(y);
	if(x!=y)
	{
		if(fa[x]>fa[y]){fa[y]+=fa[x];fa[x]=y;}
		else {fa[x]+=fa[y];fa[y]=x;}
	}
}
bool judge(int x, int y)
{
    int flag1=find(x),flag2=find(x+140000),flag3=find(y),flag4=find(y+140000);
    if(flag1==flag3)return 1;
    if(flag2==flag4)return 1;
    connect(flag1,flag4);
    connect(flag2,flag3);
    return 0;
}
int main()
{
	scanf("%d%d",&n,&K);
	for(int i=1;i<=n;i++){scanf("%d",&a[i]);maxn=max(maxn,a[i]);}
	if(K==2)
	{
		for(int i=n,j=n;i;)
		{
			for(;j;j--)
			{
				if(vis[a[j]])
				{
					if((int)sqrt(a[j]<<1)*(int)sqrt(a[j]<<1)==a[j]<<1)
					{
						if(sit[a[j]])break;
						for(int k=1;k*k-a[j]<=maxn;k++)
						{
							if(k*k-a[j]<0)continue;
							if(vis[k*k-a[j]]&&k*k!=a[j]<<1)goto nxt;
						}
						sit[a[j]]=1;
					}
				}
				else
				{
					for(int k=1;k*k-a[j]<=maxn;k++)
					{
						if(k*k-a[j]<=0)continue;
						if(vis[k*k-a[j]]&&(judge(k*k-a[j],a[j])||sit[a[j]]||sit[k*k-a[j]])){fa[a[j]]=fa[a[j]+140000]=0;goto nxt;}
					}
					vis[a[j]]=1;
				}
			}
			nxt:
			if(!j)break;
			que[++sum]=j;
			for(;i>j;i--)fa[a[i]]=fa[a[i]+140000]=vis[a[i]]=sit[a[i]]=0;
		}
		printf("%d\n",sum+1);
		for(int i=sum;i;i--)printf("%d ",que[i]);
	}
	return 0;
}

$100$分代码：

#include<bits/stdc++.h>
using namespace std;
int n,K;
int a[150000],maxn;
bool vis[150000],sit[150000];
int que[150000],sum;
int fa[300000];
void solve1()
{
	for(int i=n,j=n;i;)
	{
		for(;j;j--)
		{
			for(int k=1;k*k-a[j]<=maxn;k++)
			{
				if(k*k-a[j]<=0)continue;
				if(vis[k*k-a[j]])goto nxt;
			}
			vis[a[j]]=1;
		}
		nxt:
		if(!j)break;
		que[++sum]=j;
		for(;i>j;i--)vis[a[i]]=0;
	}	
}
int find(int x){return fa[x]<=0?x:fa[x]=find(fa[x]);}
void connect(int x,int y)
{
	x=find(x);
	y=find(y);
	if(x!=y)
	{
		if(fa[x]>fa[y]){fa[y]+=fa[x];fa[x]=y;}
		else {fa[x]+=fa[y];fa[y]=x;}
	}
}
bool judge(int x, int y)
{
    int flag1=find(x),flag2=find(x+140000),flag3=find(y),flag4=find(y+140000);
    if(flag1==flag3)return 1;
    if(flag2==flag4)return 1;
    connect(flag1,flag4);
    connect(flag2,flag3);
    return 0;
}
void solve2()
{
	for(int i=n,j=n;i;)
	{
		for(;j;j--)
		{
			if(vis[a[j]])
			{
				if((int)sqrt(a[j]<<1)*(int)sqrt(a[j]<<1)==a[j]<<1)
				{
					if(sit[a[j]])break;
					for(int k=1;k*k-a[j]<=maxn;k++)
					{
						if(k*k-a[j]<0)continue;
						if(vis[k*k-a[j]]&&k*k!=a[j]<<1)goto nxt;
					}
					sit[a[j]]=1;
				}
			}
			else
			{
				for(int k=1;k*k-a[j]<=maxn;k++)
				{
					if(k*k-a[j]<=0)continue;
					if(vis[k*k-a[j]]&&(judge(k*k-a[j],a[j])||sit[a[j]]||sit[k*k-a[j]])){fa[a[j]]=fa[a[j]+140000]=0;goto nxt;}
				}
				vis[a[j]]=1;
			}
		}
		nxt:
		if(!j)break;
		que[++sum]=j;
		for(;i>j;i--)fa[a[i]]=fa[a[i]+140000]=vis[a[i]]=sit[a[i]]=0;
	}
}		
int main()
{
	scanf("%d%d",&n,&K);
	for(int i=1;i<=n;i++){scanf("%d",&a[i]);maxn=max(maxn,a[i]);}
	if(K==1)solve1();
	else solve2();
	printf("%d\n",sum+1);
	for(int i=sum;i;i--)printf("%d ",que[i]);
	return 0;
}

---

rp++