Kd树寻找距离P最近的K个点（C++ 代码实现）

推荐学习该知识文章

实习时划水手动实现代码：

#include <algorithm>
#include <cstdlib>
#include <iostream>
#include <queue>
#include <unordered_map>
#include <vector>
using namespace std;
struct KdNode {
  KdNode* Parent;
  KdNode* LSon;
  KdNode* RSon;
  bool div;
  float Axis, Axis_x, Axis_y;
  KdNode() { Parent = nullptr, LSon = nullptr, RSon = nullptr; }
  KdNode(KdNode* Par, KdNode* Ls, KdNode* Rs, float axis, float axis_x,
         float axis_y)
      : Parent(Par),
        LSon(Ls),
        RSon(Rs),
        Axis(axis),
        Axis_x(axis_x),
        Axis_y(axis_y) {}
};
unordered_map<KdNode*, bool> vis;
priority_queue<pair<float, KdNode*>, vector<pair<float, KdNode*> > > q;
struct Point {
  float Posx, Posy;  //坐标轴x
} KdAxis[1005];

bool cmpx(Point Pxo, Point Pxt) { return Pxo.Posx < Pxt.Posx; }

bool cmpy(Point Pxo, Point Pxt) { return Pxo.Posy < Pxt.Posy; }

inline float GetDis(float Kdx, float Kdy, float Px, float Py) {
  return (Kdx - Px) * (Kdx - Px) + (Kdy - Py) * (Kdy - Py);
}

void BuildKdTree(KdNode* Rt, KdNode* Fa, int L, int R, int div) {
  // cout << L << '*' << R << endl;
  Rt->Parent = Fa;
  Rt->div = div;
  if (!div) {
    sort(KdAxis + L, KdAxis + R + 1, cmpx);
  } else {
    sort(KdAxis + L, KdAxis + R + 1, cmpy);
  }

  div = (div + 1) % 2;
  int MID = (L + R) / 2;
  KdNode* Ls = new KdNode();
  KdNode* Rs = new KdNode();

  Rt->Axis = div == 0 ? KdAxis[MID].Posx : KdAxis[MID].Posy;
  Rt->Axis_x = KdAxis[MID].Posx;
  Rt->Axis_y = KdAxis[MID].Posy;
  if (L <= MID - 1) {
    Rt->LSon = Ls;
  }
  if (MID + 1 <= R) {
    Rt->RSon = Rs;
  }
  if (L <= MID - 1) {
    BuildKdTree(Ls, Rt, L, MID - 1, div);
  }
  if (MID + 1 <= R) {
    BuildKdTree(Rs, Rt, MID + 1, R, div);
  }
}

void dfs(KdNode* Rt, float px, float py, int k) {
  if (Rt == nullptr) {
    return;
  }
  if (!vis[Rt]) {
    if (Rt->div == 0) {
      if (px <= Rt->Axis_x) {
        dfs(Rt->LSon, px, py, k);

        if (!vis[Rt]) {
          if (q.size() < k) {
            q.push(make_pair(GetDis(Rt->Axis_x, Rt->Axis_y, px, py), Rt));
          } else {
            if (GetDis(Rt->Axis_x, Rt->Axis_y, px, py) < q.top().first) {
              q.pop();
              q.push(make_pair(GetDis(Rt->Axis_x, Rt->Axis_y, px, py), Rt));
            }
          }
        }

        vis[Rt] = true;
        if (q.top().first > GetDis(Rt->Axis, 0, px, py) || q.size() < k) {
          dfs(Rt->RSon, px, py, k);
        }

      } else {
        dfs(Rt->RSon, px, py, k);

        if (!vis[Rt]) {
          if (q.size() < k) {
            q.push(make_pair(GetDis(Rt->Axis_x, Rt->Axis_y, px, py), Rt));
          } else {
            if (GetDis(Rt->Axis_x, Rt->Axis_y, px, py) < q.top().first) {
              q.pop();
              q.push(make_pair(GetDis(Rt->Axis_x, Rt->Axis_y, px, py), Rt));
            }
          }
        }

        vis[Rt] = true;
        if (q.top().first > GetDis(Rt->Axis, 0, px, py) || q.size() < k) {
          dfs(Rt->LSon, px, py, k);
        }
      }
    } else {
      if (py <= Rt->Axis_y) {
        dfs(Rt->LSon, px, py, k);

        if (!vis[Rt]) {
          if (q.size() < k) {
            q.push(make_pair(GetDis(Rt->Axis_x, Rt->Axis_y, px, py), Rt));
          } else {
            if (GetDis(Rt->Axis_x, Rt->Axis_y, px, py) < q.top().first) {
              q.pop();
              q.push(make_pair(GetDis(Rt->Axis_x, Rt->Axis_y, px, py), Rt));
            }
          }
        }

        vis[Rt] = true;
        if (q.top().first > GetDis(0, Rt->Axis, px, py) || q.size() < k) {
          dfs(Rt->RSon, px, py, k);
        }
      } else {
        dfs(Rt->RSon, px, py, k);

        if (!vis[Rt]) {
          if (q.size() < k) {
            q.push(make_pair(GetDis(Rt->Axis_x, Rt->Axis_y, px, py), Rt));
          } else {
            if (GetDis(Rt->Axis_x, Rt->Axis_y, px, py) < q.top().first) {
              q.pop();
              q.push(make_pair(GetDis(Rt->Axis_x, Rt->Axis_y, px, py), Rt));
            }
          }
        }

        vis[Rt] = true;
        if (q.top().first > GetDis(0, Rt->Axis, px, py) || q.size() < k) {
          dfs(Rt->LSon, px, py, k);
        }
      }
    }
  }
}

int main() {
  KdAxis[1].Posx = 6.27, KdAxis[1].Posy = 5.50;
  KdAxis[2].Posx = 1.24, KdAxis[2].Posy = -2.86;
  KdAxis[3].Posx = -6.88, KdAxis[3].Posy = -5.40;
  KdAxis[4].Posx = -2.96, KdAxis[4].Posy = -2.50;
  KdAxis[5].Posx = -4.60, KdAxis[5].Posy = -10.55;
  KdAxis[6].Posx = -4.96, KdAxis[6].Posy = 12.61;
  KdAxis[7].Posx = 1.75, KdAxis[7].Posy = 12.26;
  KdAxis[8].Posx = 17.05, KdAxis[8].Posy = -12.79;
  KdAxis[9].Posx = 7.75, KdAxis[9].Posy = -22.68;
  KdAxis[10].Posx = 15.31, KdAxis[10].Posy = -13.16;
  KdAxis[11].Posx = 10.80, KdAxis[11].Posy = -5.03;
  KdAxis[12].Posx = 7.83, KdAxis[12].Posy = 15.70;
  KdAxis[13].Posx = 14.63, KdAxis[13].Posy = -0.35;

  KdNode* Rt = new KdNode();
  int div = 0, L = 1, R = 13;
  BuildKdTree(Rt, nullptr, L, R, div);

  dfs(Rt, -1, -5, 5);
  while (!q.empty()) {
    KdNode* NowNode = q.top().second;
    cout << NowNode->Axis_x << ' ' << NowNode->Axis_y << endl;
    q.pop();
  }
  return 0;
}