Luogu P3940 分组 (带权并查集动态判二分图)

题意

传送门

题解

兔子越多限制越大，所以字典序最小就肯定是最后一组尽量多。从后往前枚举加兔子，如果矛盾就分组。

中间不能出现奇环，只能是一个二分图，用带权并查集动态判。

要特判a+a=k^2的情况。

CODE

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int MAXN = (1<<18)+5;
const int SQRT = 1<<9;
int n, m, k, a[MAXN], sq[MAXN], f[MAXN];
bool d[MAXN], chk[MAXN];
int ans[MAXN], vis[MAXN];
inline void solve1() {
	for(int i = n; i; --i) {
		bool flg = 0;
		for(int j = sq[a[i]]+1; j <= SQRT; ++j)
			if(vis[j*j-a[i]]) { flg = 1; break; }
		if(flg) { for(int j = i+1; j <= ans[m]; ++j) vis[a[j]] = 0; ans[++m] = i; }
		++vis[a[i]];
	}
}
int find(int x) {
	if(x != f[x]) {
		int fa = f[x];
		f[x] = find(fa);
		d[x] ^= d[fa];
	}
	return f[x];
}
inline void solve2() {
	int x, y, k;
	for(int i = n; i; --i) {
		bool flg = 0, same = 0, other = 0;
		for(int j = sq[a[i]]+1; j <= SQRT; ++j)
			if(vis[k=j*j-a[i]]) {
				if(vis[k] > 1 && chk[k]) { flg = 1; break; }
				if(a[i] == k) { same = 1; continue; }
				other = 1;
				if((x=find(a[i])) != (y=find(k))) f[x] = y, d[x] = d[a[i]] ^ d[k] ^ 1;
				else if(d[a[i]] == d[k]) { flg = 1; break; }
			}
		if(flg || (same & other)) { for(int j = i+1; j <= ans[m]; ++j) vis[a[j]] = 0, f[a[j]] = a[j], d[a[j]] = 0; ans[++m] = i; }
		++vis[a[i]];
	}
}
int main () {
	for(int i = 1; i <= SQRT; ++i) { sq[i*i] = i; if(!(i&1)) chk[i*i/2] = 1; }
	for(int i = 1; i < MAXN; ++i) { sq[i] = sq[i] ? sq[i] : sq[i-1]; f[i] = i, d[i] = 0; }
	scanf("%d%d", &n, &k); ans[m=0] = n;
	for(int i = 1; i <= n; ++i) scanf("%d", &a[i]);
	if(k == 1) solve1();
	else solve2();
	printf("%d
", m+1);
	for(int i = m; i; --i) printf("%d%c", ans[i], " 
"[i==1]);
}