K-D树

2024-08-15

K-D树能将二维空间上的点集转化为二叉树,求解最近点问题

#include <bits/stdc++.h>
#define io cin.tie(0), cout.tie(0), ios::sync_with_stdio(false)
#define LL long long
#define ULL unsigned long long
#define EPS 1e-8
#define INF 0x7fffffff
#define SUB -INF - 1
using namespace std;
const int N = 100005;
const int K = 2;
struct Point
{
    int dim[K];
};
Point q[N];
Point t[N];
int now;
bool cmp(Point a, Point b) { return a.dim[now] < b.dim[now]; }
LL square(int x) { return (LL)x * x; }
LL dis(Point a, Point b) { return square(a.dim[0] - b.dim[0]) + square(a.dim[1] - b.dim[1]); }
void build(int L, int R, int dep)
{
    if (L >= R)
        return;
    int d = dep % K;
    int mid = (L + R) >> 1;
    now = d;
    nth_element(t + L, t + mid, t + R, cmp);
    build(L, mid, dep + 1);
    build(mid + 1, R, dep + 1);
}
LL ans;
void query(int L, int R, int dep, Point p)
{
    if (L >= R)
        return;
    int mid = (L + R) >> 1;
    int d = dep % K;
    LL mindis = dis(t[mid], p);
    if (ans == 0)
        ans = mindis;
    if (mindis != 0 && ans > mindis)
        ans = mindis;
    if (p.dim[d] > t[mid].dim[d])
    {
        query(mid + 1, R, dep + 1, p);
        if (ans > square(t[mid].dim[d] - p.dim[d]))
            query(L, mid, dep + 1, p);
    }
    else
    {
        query(L, mid, dep + 1, p);
        if (ans > square(t[mid].dim[d] - p.dim[d]))
            query(mid + 1, R, dep + 1, p);
    }
}
int main()
{
    int T;
    scanf("%d", &T);
    while (T--)
    {
        int n;
        scanf("%d", &n);
        for (int i = 0; i < n; i++)
            scanf("%d%d", &(q[i].dim[0]), &(q[i].dim[1])), t[i] = q[i];
        build(0, n, 0);
        for (int i = 0; i < n; i++)
        {
            ans = 0;
            query(0, n, 0, q[i]);
            printf("%lld\n", ans);
        }
    }
    return 0;
}