【BZOJ4520】【CQOI2016】K远点对

【题目链接】

• 点击打开链接
  

【思路要点】

• KDTree实现K-临近搜索，实现时需要用一个堆来辅助。
• 时间复杂度\(O(N\sqrt{N}+NKLogK)\)。

【代码】

#include<bits/stdc++.h>
using namespace std;
#define MAXN	100005
template <typename T> void read(T &x) {
	x = 0; int f = 1;
	char c = getchar();
	for (; !isdigit(c); c = getchar()) if (c == '-') f = -f;
	for (; isdigit(c); c = getchar()) x = x * 10 + c - '0';
	x *= f;
}
struct point {long long x, y; };
int cmptype;
bool cmp(point a, point b) {
	if (cmptype) return (a.x < b.x) || (a.x == b.x && a.y < b.y);
	else return (a.y < b.y) || (a.y == b.y && a.x < b.x);
}
point Min(point a, point b) {
	a.x = min(a.x, b.x);
	a.y = min(a.y, b.y);
	return a;
}
point Max(point a, point b) {
	a.x = max(a.x, b.x);
	a.y = max(a.y, b.y);
	return a;
}
struct KD_Tree {
	struct Node {
		point pos;
		point leftup, rightdown;
		int lc, rc;
	};
	Node a[MAXN];
	point p[MAXN], now;
	int n, k, root, size;
	priority_queue <long long, vector<long long>, greater<long long> > Heap;
	void update(int root) {
		a[root].leftup = a[root].rightdown = a[root].pos;
		if (a[root].lc) {
			int tmp = a[root].lc;
			a[root].leftup = Min(a[root].leftup, a[tmp].leftup);
			a[root].rightdown = Max(a[root].rightdown, a[tmp].rightdown);
		}
		if (a[root].rc) {
			int tmp = a[root].rc;
			a[root].leftup = Min(a[root].leftup, a[tmp].leftup);
			a[root].rightdown = Max(a[root].rightdown, a[tmp].rightdown);
		}
	}
	void build(int l, int r, int &root, int type) {
		if (root == 0) root = ++size;
		int mid = (l + r) / 2;
		cmptype = type;
		nth_element(p + l, p + mid, p + r + 1, cmp);
		a[root].pos = p[mid];
		if (mid > l) build(l, mid - 1, a[root].lc, type ^ 1);
		if (mid < r) build(mid + 1, r, a[root].rc, type ^ 1);
		update(root);
	}
	void init() {
		read(n); read(k);
		size = 0;
		for (int i = 1; i <= n; i++)
			read(p[i].x), read(p[i].y);
		build(1, n, root, 0);
	}
	long long dist(point a, point b) {
		return (a.x - b.x) * (a.x - b.x) + (a.y - b.y) * (a.y - b.y);
	}
	long long dist(point now, int root) {
		long long tmp = max(dist(now, a[root].leftup), dist(now, (point) {a[root].leftup.x, a[root].rightdown.y}));
		long long tnp = max(dist(now, a[root].rightdown), dist(now, (point) {a[root].rightdown.x, a[root].leftup.y}));
		return max(tmp, tnp);
	}
	void work(int root) {
		if (Heap.size() == 2 * k && dist(now, root) <= Heap.top()) return;
		long long tmp = dist(now, a[root].pos);
		Heap.push(tmp);
		if (Heap.size() > 2 * k) Heap.pop();
		if (a[root].lc) work(a[root].lc);
		if (a[root].rc) work(a[root].rc);
	}
	long long calc() {
		for (int i = 1; i <= n; i++) {
			now = p[i];
			work(root);
		}
		return Heap.top();
	}
} KDT;
int main() {
	KDT.init();
	cout << KDT.calc() << endl;
	return 0;
}