【BZOJ 3053】The Closest M Points

时间：2021-09-16 02:34:13

KDTree模板，在m维空间中找最近的k个点，用的是欧几里德距离。

理解了好久，昨晚始终不明白那些“估价函数”，后来才知道分情况讨论，≤k还是=k，在当前这一维度距离过线还是不过线，过线则要继续搜索另一个子树。还有别忘了当前这个节点！

#include<cmath>

#include<queue>

#include<cstdio>

#include<cstring>

#include<algorithm>

#define read(x) x=getint()

using namespace std;

typedef long long LL;

const int N = 50003;

const int inf = 0x7fffffff;

int getint() {

	int k = 0, fh = 1; char c = getchar();

	for(; c < '0' || c > '9'; c = getchar())

		if (c == '-') fh = -1;

	for(; c >= '0' && c <= '9'; c = getchar())

		k = k * 10 + c - '0';

	return k * fh;

}

int n, m, root, D;

LL minn;

struct P {

	int d[5], mx[5], mn[5], l, r, id;

	P (): l(0), r(0), id(0) {};

	int &operator [] (int x) {return d[x];}

	bool operator < (P point) const {return d[D] < point[D];}

} T[N << 1], po[N];

priority_queue <pair <LL, int> > Q;

LL sqr(LL x) {return x * x;}

LL dis(P a, P b) {

	LL ret = 0;

	for(int i = 0; i < m; ++i) ret += sqr(a[i] - b[i]);

	return ret;

}

void pushup(int x, int y) {

	for(int i = 0; i < m; ++i)

		T[x].mn[i] = min(T[x].mn[i], T[y].mn[i]),

		T[x].mx[i] = max(T[x].mx[i], T[y].mx[i]);

}

int Build(int l, int r, int dd) {

	D = dd; int mid = (l + r) >> 1; nth_element(po + l, po + mid, po + r + 1);

	for(int i = 0; i < m; ++i)

		T[mid].mn[i] = T[mid].mx[i] = T[mid].d[i] = po[mid].d[i];

	T[mid].id = mid;

	if (l < mid) T[mid].l = Build(l, mid - 1, (dd + 1) % m);

	if (mid < r) T[mid].r = Build(mid + 1, r, (dd + 1) % m);

	if (T[mid].l) pushup(mid, T[mid].l);

	if (T[mid].r) pushup(mid, T[mid].r);

	return mid;

}

void ask(int rt, int dd, P p, int k) {

	int L = T[rt].l, R = T[rt].r;

	if (p[dd] >= T[rt][dd]) swap(L, R);

	if (L) ask(L, (dd + 1) % m, p, k);

	bool pd = 0; LL di = dis(T[rt], p); minn = min(minn, di);

	if (Q.size() < k) { Q.push(pair <LL, int> (di, rt)); pd = 1;}

	else {

		if (di < Q.top().first) Q.pop(), Q.push(make_pair(di, rt));

		if (sqr(p[dd] - T[rt][dd]) < Q.top().first) pd = 1;

	}

	if (pd && R) ask(R, (dd + 1) % m, p, k);

}

int ans[N];

int main() {

	while (~scanf("%d%d", &n, &m)) {

		for(int i = 1; i <= n; ++i)

			for(int j = 0; j < m; ++j)

				read(po[i][j]);

		memset(T, 0, sizeof(T));

		root = Build(1, n, 0);

		int qq; read(qq);

		for(; qq; --qq) {

			P p; int k;

			for(int i = 0; i < m; ++i)

				read(p[i]);

			read(k);

			printf("the closest %d points are:\n", k);

			minn = inf; ask(root, 0, p, k);

			while (!Q.empty()) {

				ans[++ans[0]] = Q.top().second;

				Q.pop();

			}

			for(; ans[0]; --ans[0])

				for(int i = 0; i < m; ++i)

					printf("%d%c", T[ans[ans[0]]][i]," \n"[i == m - 1]);

		}

	}

	return 0;

}

我就是弱啊～～～

相关文章