P3958 [NOIP2017 提高组] 奶酪（并查集）

[NOIP2017 提高组] 奶酪 - 洛谷https://www.luogu.com.cn/problem/P3958

#include <iostream>
#include <cstdio>
#include <string>
#include <algorithm>
#include <vector>
#include <queue>
#include <stack>
#include <cstring>
#include <set>
#include <cmath>
#include <map>
#include <cstdlib>
typedef long long ll;
typedef unsigned long long ull;
using namespace std;
const int MN = 65005;
const int MAXN = 2000010;
const int INF = 0x3f3f3f3f;
#define IOS ios::sync_with_stdio(false)

int pre[MAXN];

int find(int x) {
	if (pre[x] == x) {
		return x;
	}
	return pre[x] = find(pre[x]);
}

ll dis(ll x1, ll y1, ll z1, ll x2, ll y2, ll z2) {
	return (x1 - x2) * (x1 - x2) + (y1 - y2) * (y1 - y2) + (z1 - z2) * (z1 - z2);
}
ll x[10005], y[10005], z[10005];
int f1[10001], f2[100005];

inline bool same(int x, int y) {
	return find(x) == find(y);
}

int main() {
	int t;
	scanf("%d", &t);
	int n, h;
	ll r;
	while (t--) {
		scanf("%d %d %lld", &n, &h, &r);
		int cnt1 = 0;
		int cnt2 = 0;
		for (int i = 1; i <= n; i++) {
			pre[i] = i;
		}
		for (int i = 1; i <= n; i++) {
			scanf("%lld %lld %lld", x + i, y + i, z + i);
			if (z[i] + r >= h) {
				cnt1++;
				f1[cnt1] = i;
			}
			if (z[i] - r <= 0) {
				cnt2++;
				f2[cnt2] = i;
			}
			for (int j = 1; j < i; j++) {
				if (dis(x[i], y[i], z[i], x[j], y[j], z[j]) <= 4 * r * r) {
					pre[find(i)] = find(j);
				}
			}
		}
		bool flag=true;
		for (int i = 1; i <= cnt1; i++) {
			for (int j = 1; j <= cnt2; j++) {
				if (same(f1[i], f2[j])) {
					printf("Yes\n");
					flag =false;
					break;
				}
			}
			if (!flag) {
				break;
			}
		}
		if(flag){
			printf("No\n");
		}
	}
	return 0;
}