K-D Tree

kd_tree.cpp

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using ull = unsigned long long;

struct point {
	int x, y;
	point operator-(point v) const { return { x - v.x, y - v.y }; }
	ll operator~() const { return ll(x) * x + ll(y) * y; }
};

template<int LeafSize = 1>
struct kd_tree : vector<point> {
	#define A (*this)
	point p;
	ull dist;
	kd_tree(vector<point> const& _): vector(_) {
		if(size() > LeafSize) qry<0, 1>(0, size());
	}
	template<bool coord = 0, bool is_build = 0>
	void qry(int s, int e) {
		if(s >= e) return;
		if(e - s <= LeafSize) {
			for(int i = s; i < e; i++)
				dist = min(dist, ull(~(A[i] - p) - 1));
			return;
		}
		int m = (s + e) >> 1;
		if constexpr(is_build) {
			nth_element(begin() + s, begin() + m, begin() + e, [](point l, point r) { return coord ? l.x < r.x : l.y < r.y; });
			if(m - s > LeafSize) qry<!coord, 1>(s, m);
			if(e - m - 1 > LeafSize) qry<!coord, 1>(m + 1, e);
		} else {
			dist = min(dist, ull(~(A[m] - p) - 1));
			int c = (coord ? p.x - A[m].x : p.y - A[m].y);
			if(c < 0) {
				qry<!coord>(s, m);
				if(ll(c) * c < dist + 1) qry<!coord>(m + 1, e);
			} else {
				qry<!coord>(m + 1, e);
				if(ll(c) * c < dist + 1) qry<!coord>(s, m);
			}
		}
	}
	auto qry(point q) {
		p = q;
		dist = ULLONG_MAX;
		qry(0, size());
		return dist + 1;
	}
	#undef A
};

int main() {
	ios::sync_with_stdio(0);cin.tie(0);
	int TC;
	cin >> TC;
	while(TC--) {
		int N;
		cin >> N;
		vector<point> A(N);
		for(auto& [ x, y ] : A)
			cin >> x >> y;

		kd_tree T(A);
		for(auto p : A)
			cout << T.qry(p) << '\n';
	}
}

kd_tree_etc.cpp

#pragma GCC optimize("O3")
#pragma GCC target("avx2")
#define __SSE4_1__
#define __AVX2__
#include <bits/stdc++.h>
#include <experimental/simd>
using namespace std;
using namespace experimental;
using ll = long long;
using ull = unsigned long long;
template<typename T> using sv = native_simd<T>;
struct point { int x, y; };

int N;
point A[100000], B[100000], p;
alignas(memory_alignment_v<sv<int>>) int X[100000], Y[100000];
ull dist;

void upd(int i) {
	dist = min(dist, ull((X[i] - p.x)*ll(X[i] - p.x) + (Y[i] - p.y)*ll(Y[i] - p.y) - 1));
}

constexpr int LeafSize = 172;
template<bool coord = 0, bool is_build = 0>
void qry(int s, int e) {
	if(s >= e) return;
	if(e - s <= LeafSize) {
		int i = s;
		sv<ull> d = dist;
		for(; i & (d.size() * 2 - 1); i++) upd(i);
		for(; i + d.size() * 2 <= e; i += d.size() * 2) {
			auto xl = bit_cast<sv<ull>>(sv<int>(X + i, vector_aligned) - p.x);
			auto yl = bit_cast<sv<ull>>(sv<int>(Y + i, vector_aligned) - p.y);
			auto xh = xl >> 32;
			auto yh = yl >> 32;
			asm("vpmuldq %1, %2, %0" : "=x"(xl) : "x"(xl), "x"(xl));
			asm("vpmuldq %1, %2, %0" : "=x"(yl) : "x"(yl), "x"(yl));
			asm("vpmuldq %1, %2, %0" : "=x"(xh) : "x"(xh), "x"(xh));
			asm("vpmuldq %1, %2, %0" : "=x"(yh) : "x"(yh), "x"(yh));
			d = min(d, min(xl + yl - 1, xh + yh - 1));
		}
		dist = min(dist, hmin(d));
		for(; i < e; i++) upd(i);
		return;
	}
	int m = (s + e) >> 1;
	if constexpr(is_build) {
		nth_element(A + s, A + m, A + e, [](point l, point r) { return coord ? l.x < r.x : l.y < r.y; });
		if(m - s > LeafSize) qry<!coord, 1>(s, m);
		if(e - m - 1 > LeafSize) qry<!coord, 1>(m + 1, e);
	} else {
		upd(m);
		int c = (coord ? p.x - A[m].x : p.y - A[m].y);
		if(c < 0) {
			qry<!coord>(s, m);
			if(ll(c) * c < dist + 1) qry<!coord>(m + 1, e);
		} else {
			qry<!coord>(m + 1, e);
			if(ll(c) * c < dist + 1) qry<!coord>(s, m);
		}
	}
}
auto qry(point q) {
	p = q;
	dist = ULLONG_MAX;
	qry(0, N);
	return dist + 1;
}

int main() {
	ios::sync_with_stdio(0);cin.tie(0);
	int TC;
	cin >> TC;
	while(TC--) {
		cin >> N;
		for(int i = 0; i < N; i++)
			cin >> A[i].x >> A[i].y;
		copy(A, A + N, B);
		if(N > LeafSize) qry<0, 1>(0, N);
		for(int i = 0; i < N; i++)
			X[i] = A[i].x, Y[i] = A[i].y;
		for(int i = 0; i < N; i++)
			cout << qry(B[i]) << '\n';
	}
}