KdTreeを自作してみる（２）探索

class mykdtree {

   std::vector<Eigen::Vector3d>& _points;
   node* node_root;
public:
   mykdtree(std::vector<Eigen::Vector3d>& points) : _points(points) {}

   // ツリーの構築
   void build() {

      // 最初のindexリストを作成
      std::vector<size_t> indices(_points.size());
      for (size_t i = 0; i < _points.size(); i++) {
         indices[i] = i;
      }

      // 再帰的にツリー作成
      node_root = buildnode(_points, indices, 0/*X軸*/);

   }
   node* getRoot() {
      return node_root;
   }

   // 半径内の頂点を探索する関数
   // 再帰的に探索する関数
   // nd: 現在のノード
   // query: 探索する点
   // radius: 探索半径
   // results: 探索結果のindexリスト
   void radiusSearch_core(
      node* nd,
      const Eigen::Vector3d& query,
      double radius,
      std::vector<size_t>& results) {

      if (nd == nullptr) {
         return;
      }

      const auto& nodePt = _points[nd->index];

      // ノード と query の距離を計算
      float dist = (query - nodePt).norm();


      // 現在のノードがqueryを中心とした探索範囲内にあるか確認
      if (dist <= radius) {
         results.push_back(nd->index);//現在のノードを結果に追加
      }


      // サブツリーの探索条件
      // 軸に沿った距離を計算
      ///////////////////////////////////////////////////////////
      // 
      //                            |
      // ------------●-------------○----------------
      //           query          nodePt
      //                            |
      // 
      //         query.x - nodePt.x < 0
      //           queryが現在のノードの左側にある → nodePtの左側を探索
      ///////////////////////////////////////////////////////////
      // 
      //             |
      // ------------○-------------●----------------
      //           nodePt          query
      //             |
      // 
      //         query.x - nodePt.x > 0
      //           queryが現在のノードの右側にある → nodePtの右側を探索
      ///////////////////////////////////////////////////////////

      float axisDist = query[nd->axis] - nodePt[nd->axis];

      // ケース１　queryがノードの左側にある
      if (axisDist < 0) {

         // 左側にあるので、左側を探索
         radiusSearch_core(nd->left, query, radius, results);

         // ケース１（２）nodePtが、queryの半径以内にあるなら、右側にもあるかもしれない
         if (axisDist * axisDist <= radius * radius) {
            radiusSearch_core(nd->right, query, radius, results);
         }
      }
      else {
         // ケース２　queryがノードの右側にある

         // 右側にあるので、右側を探索
         radiusSearch_core(nd->right, query, radius, results);

         // ケース２（２）nodePtが、queryの半径以内にあるなら、左側にもあるかもしれない
         if (axisDist * axisDist <= radius * radius) {
            radiusSearch_core(nd->left, query, radius, results);
         }
      }
   }

   // 半径内の頂点を探索する関数
   // query: 探索する点
   // radius: 探索半径
   // results: 探索結果のindexリスト
   void radiusSearch(
      const Eigen::Vector3d& query,
      double radius,
      std::vector<size_t>& results) {

      results.clear();
      radiusSearch_core(node_root, query, radius, results);
   }

};

#include <iostream>

//VTK_MODULE_INITに必要
#include <vtkAutoInit.h>


#include <vtkSmartPointer.h>
#include <vtkRenderer.h>
#include <vtkRenderWindow.h>
#include <vtkRenderWindowInteractor.h>

//円筒とその表示に必要
#include <vtkPolyDataMapper.h>

#pragma comment(lib,"opengl32.lib")
#pragma comment(lib,"psapi.lib")
#pragma comment(lib,"dbghelp.lib")
#pragma comment(lib,"ws2_32.lib")

#include <Eigen/Core>
#include <array>

#include <vtkActor.h>
#include <vtkPoints.h>
#include <vtkPolyData.h>
#include <vtkUnsignedCharArray.h>
#include <vtkPointData.h>
#include <vtkVertexGlyphFilter.h>
#include <vtkProperty.h>

#include "mykdtree.hpp"

//必須
VTK_MODULE_INIT(vtkRenderingOpenGL2);
VTK_MODULE_INIT(vtkInteractionStyle);

// VTK表示用
struct MyVtkCloud {
   std::vector<Eigen::Vector3d> points;
   std::array<unsigned char, 3> color;
   std::vector<std::array<unsigned char, 3> > color_array;

   vtkSmartPointer<vtkActor> actor;
   void makeActor() {
      // VTKのデータ構造に変換
      vtkSmartPointer<vtkPoints> vtk_points = vtkSmartPointer<vtkPoints>::New();
      vtkSmartPointer<vtkUnsignedCharArray> vtk_colors = vtkSmartPointer<vtkUnsignedCharArray>::New();
      vtk_colors->SetNumberOfComponents(3); // RGB
      vtk_colors->SetName("Colors");

      for (size_t i = 0; i < points.size(); ++i) {
         // 点を追加
         vtk_points->InsertNextPoint(points[i].x(), points[i].y(), points[i].z());

         if (color_array.size() == 0) {
            vtk_colors->InsertNextTypedTuple(color.data());
         }
         else {
            vtk_colors->InsertNextTypedTuple(color_array[i].data());
         }
      }

      ////////////////////////////////////////////////////////////
      vtkSmartPointer<vtkPolyData> polyData = vtkSmartPointer<vtkPolyData>::New();
      polyData->SetPoints(vtk_points);
      polyData->GetPointData()->SetScalars(vtk_colors);
      ////////////////////////////////////////////////////////////
      vtkSmartPointer<vtkVertexGlyphFilter> vertexFilter = vtkSmartPointer<vtkVertexGlyphFilter>::New();
      vertexFilter->SetInputData(polyData);
      vertexFilter->Update();
      ////////////////////////////////////////////////////////////
      vtkSmartPointer<vtkPolyDataMapper> mapper = vtkSmartPointer<vtkPolyDataMapper>::New();
      mapper->SetInputConnection(vertexFilter->GetOutputPort());

      vtkSmartPointer<vtkActor> actor = vtkSmartPointer<vtkActor>::New();
      actor->SetMapper(mapper);

      // 頂点サイズを指定
      actor->GetProperty()->SetPointSize(5); // ここで頂点サイズを指定します

      this->actor = actor;
   }
};

int main(int /*argc*/, char** /*argv*/)
{
   MyVtkCloud cloud1;

   std::vector<std::array<unsigned char, 3> > color_array;


   srand((unsigned int)5);
   //// ランダムな点群を作成
   for (size_t i = 0; i < 10000; ++i) {
      cloud1.points.push_back(Eigen::Vector3d::Random());

      color_array.push_back({ 255,0,0 });
   }
   cloud1.color = std::array<unsigned char, 3>{ 0, 255, 0 };
   cloud1.makeActor();


   mykdtree kdtree(cloud1.points);
   kdtree.build();
   node* root = kdtree.getRoot();

   std::vector<size_t> indices;
   kdtree.radiusSearch(
      Eigen::Vector3d(0.3,0.3,0.3), 0.3, indices);

   for(auto i:indices){
      color_array[i] = { 0,255,0 };
   }
   cloud1.color_array = color_array;
   cloud1.makeActor();

   // 表示
   std::vector<size_t> indices_look = indices;
   // 昇順ソート
   std::sort(indices_look.begin(), indices_look.end());
   // コンソールに表示
   for (auto i : indices_look) {
      std::cout << i << std::endl;
   }


   //////////////////////////////////////
   auto renderer = vtkSmartPointer<vtkRenderer>::New();
   renderer->AddActor(cloud1.actor);
   renderer->ResetCamera();

   //////////////////////////////////////
   auto interactor = vtkSmartPointer<vtkRenderWindowInteractor>::New();

   //////////////////////////////////////
   auto renderWindow = vtkSmartPointer<vtkRenderWindow>::New();
   renderWindow->AddRenderer(renderer);
   renderWindow->SetInteractor(interactor);
   renderWindow->Render();

   interactor->Start(); //イベントループへ入る

   return 0;
}