K-D Tree 学习笔记

\(\text{K-D Tree 学习笔记}\)

K-D Tree 是处理 \(k\) 维空间的数据结构。具有二叉搜索树的形态。事实上 OI 中的应用就是 2-D Tree，常用来维护二维平面上一些点对的信息。

构建

要让 K-D Tree 尽量贴近一个二叉搜索树的形态，那就要让某一维度上的中位数作为根节点。对于维度的选择，有两种常见的方式：一是两个维度交替建树、二是对 \(x,y\) 分别求方差，按照方差较大的维度建树。笔者一般采用第二种方式。这样一来，每个点事实上对应了平面上一个矩形。

插入

考虑维护平衡树的方式是难以运用到 KD-Tree 上的。于是常见的方法是二进制分组或替罪羊树重构。处于容易理解的角度，笔者一般使用替罪羊树重构的方式实现。具体地，设置一个平衡因子 \(\alpha\)，当 \(\max(\operatorname{siz(lc),siz(rc)})>\alpha\times\operatorname{siz(p)}\) 时暴力重构子树。值得留意的是这样做复杂度存在退化的可能性，但是 KD-Tree 更多时候作为一种骗分的方法存在，也不太会有人会卡这个。

这里给出较为完整的模板实现：

namespace KDT {
	const db shit = 0.75;
	int tot;
	struct Node {
		int lc, rc;
		int x, y;
		int l, r, u, d;
		int sz, tp;
	} e[N * M];
	#define lc(i) e[i].lc
	#define rc(i) e[i].rc
	#define x(i) e[i].x
	#define y(i) e[i].y
	#define l(i) e[i].l
	#define r(i) e[i].r
	#define u(i) e[i].u
	#define d(i) e[i].d
	#define sz(i) e[i].sz
	#define tp(i) e[i].tp
	bool cmpx(int x, int y) {
		return x(x) < x(y);
	}
	bool cmpy(int x, int y) {
		return y(x) < y(y);
	}
	void push_up(int p) {
		sz(p) = sz(lc(p)) + sz(rc(p)) + 1;
		l(p) = r(p) = x(p);
		u(p) = d(p) = y(p);
		if (lc(p)) {
			l(p) = min(l(p), l(lc(p)));
			r(p) = max(r(p), r(lc(p)));
			u(p) = min(u(p), u(lc(p)));
			d(p) = max(d(p), d(lc(p)));
		}
		if (rc(p)) {
			l(p) = min(l(p), l(rc(p)));
			r(p) = max(r(p), r(rc(p)));
			u(p) = min(u(p), u(rc(p)));
			d(p) = max(d(p), d(rc(p)));
		}
	}
	bool gotoshit(int p) {
		return (db)max(sz(lc(p)), sz(rc(p))) > sz(p) * shit;
	}
	int id[N], cp;
	void sve(int p) {
		if (!p) return;
		sve(lc(p));
		id[++cp] = p;
		sve(rc(p));
	}
	void build(int &p, int l, int r) {
		if (l > r) return p = 0, void();
		int mid = (l + r) >> 1;
		db v1 = 0, v2 = 0, av1 = 0, av2 = 0;
		for (int i = l; i <= r; i++) {
			av1 += x(id[i]);
			av2 += y(id[i]);
		}
		av1 /= r - l + 1, av2 /= r - l + 1;
		for (int i = l; i <= r; i++) {
			v1 += (av1 - x(id[i])) * (av1 - x(id[i]));
			v2 += (av2 - y(id[i])) * (av2 - y(id[i]));
		}
		if (v1 > v2) {
			nth_element(id + l, id + mid, id + r + 1, cmpx);
			tp(p) = 1;
		}
		else {
			nth_element(id + l, id + mid, id + r + 1, cmpy);
			tp(p) = 2;
		}
		p = id[mid];
		build(lc(p), l, mid - 1);
		build(rc(p), mid + 1, r);
		push_up(p);
	}
	void forashit(int &p) {
		cp = 0;
		sve(p);
		build(p, 1, cp);
	}
	void insert(int &p, int x, int y) {
		if (!p) {
			p = ++tot;
			x(p) = x, y(p) = y;
			return push_up(p);
		}
		if (tp(p) == 1) {
			if (x <= x(p)) insert(lc(p), x, y);
			else insert(rc(p), x, y);
		}
		else {
			if (y <= y(p)) insert(lc(p), x, y);
			else insert(rc(p), x, y);
		}
		push_up(p);
		if (gotoshit(p)) forashit(p);
	}
}

查询

可以证明 K-D Tree 查询矩形范围内维护的问题复杂度是单次 \(O(n^{1-\frac 1k})\)。但是有时我们会用它求解"平面上与其最近的点"一类问题。这样的查询最坏的复杂度是 \(O(n)\)，但是使用启发式搜索+估价函数期望下是 \(O(\log n)\) 的。具体的实现是好理解的，可见代码。

int fmx(int p, int x, int y) {
	int ans = 0;
	ans += max(abs(x - l(p)), abs(x - r(p)));
	ans += max(abs(y - u(p)), abs(y - d(p)));
	return ans;
}
int fmn(int p, int x, int y) {
	int ans = 0;
	if (x < l(p)) ans += l(p) - x;
	if (x > r(p)) ans += x - r(p);
	if (y < u(p)) ans += u(p) - y;
	if (y > d(p)) ans += y - d(p);
	return ans;
}
int amx, amn;
void qmx(int p, int x, int y) {
	if (!p) return;
	amx = max(amx, abs(x - x(p)) + abs(y - y(p)));
	int vl = -inf, vr = -inf;
	if (lc(p)) vl = fmx(lc(p), x, y);
	if (rc(p)) vr = fmx(rc(p), x, y);
	if (vl > vr) {
		if (vl > amx) qmx(lc(p), x, y);
		if (vr > amx) qmx(rc(p), x, y);
	}
	else {
		if (vr > amx) qmx(rc(p), x, y);
		if (vl > amx) qmx(lc(p), x, y);
	}
}
void qmn(int p, int x, int y) {
	if (!p) return;
	if (!(x(p) == x && y(p) == y)) amn = min(amn, abs(x - x(p)) + abs(y - y(p)));
	int vl = inf, vr = inf;
	if (lc(p)) vl = fmn(lc(p), x, y);
	if (rc(p)) vr = fmn(rc(p), x, y);
	if (vl < vr) {
		if (vl < amn) qmn(lc(p), x, y);
		if (vr < amn) qmn(rc(p), x, y);
	}
	else {
		if (vr < amn) qmn(rc(p), x, y);
		if (vl < amn) qmn(lc(p), x, y);
	}
}