Apollo | common | math | kd-tree

0.引言1.建树过程2.查询

0.引言

Apollo中的KD-Tree实现代码.

KD-Tree参数：

struct AABoxKDTreeParams {
  /// The maximum depth of the kdtree.
  int max_depth = -1;
  /// The maximum number of items in one leaf node.
  int max_leaf_size = -1;
  /// The maximum dimension size of leaf node.
  // 可以看做是拆分到最后时像素间的最大距离
  double max_leaf_dimension = -1.0;
};

对外接口，模板类：

template 
class AABoxKDTree2d {
 public:
  using ObjectPtr = const ObjectType *;

  
  AABoxKDTree2d(const std::vector &objects,
                const AABoxKDTreeParams ¶ms) {
    if (!objects.empty()) {
      std::vector object_ptrs;
      for (const auto &object : objects) {
        object_ptrs.push_back(&object);
      }
      // 具体建树接口
      root_.reset(new AABoxKDTree2dNode(object_ptrs, params, 0));
    }
  }

  
  ObjectPtr GetNearestObject(const Vec2d &point) const {
    return root_ == nullptr ? nullptr : root_->GetNearestObject(point);
  }

  
  std::vector GetObjects(const Vec2d &point,
                                    const double distance) const {
    if (root_ == nullptr) {
      return {};
    }
    return root_->GetObjects(point, distance);
  }

  
  AABox2d GetBoundingBox() const {
    return root_ == nullptr ? AABox2d() : root_->GetBoundingBox();
  }

 private:
  std::unique_ptr> root_ = nullptr;
};

1.建树过程

  
  AABoxKDTree2dNode(const std::vector &objects,
                    const AABoxKDTreeParams ¶ms, int depth)
      : depth_(depth) {
    CHECK(!objects.empty());
    //Step 1.计算 Dx 边界
    ComputeBoundary(objects);
    //Step 2.对当前Node计算划分点 partition_position_
    ComputePartition();
    //Step 3.分裂当前Node为左右子树
    if (SplitToSubNodes(objects, params)) {//检查是否分裂是否满足参数条件
      std::vector left_subnode_objects;
      std::vector right_subnode_objects;
      //Step 3.1.分裂Node
      PartitionObjects(objects, &left_subnode_objects, &right_subnode_objects);

      //Step 3.2. Split to sub-nodes. 分裂到左右子树
      if (!left_subnode_objects.empty()) {
        left_subnode_.reset(new AABoxKDTree2dNode(//在构造函数中递归
            left_subnode_objects, params, depth + 1));
      }
      if (!right_subnode_objects.empty()) {
        right_subnode_.reset(new AABoxKDTree2dNode(
            right_subnode_objects, params, depth + 1));
      }
    } else {
      InitObjects(objects);//?这句会运行吗
    }
  }

step1 计算 Dx 边界

  void ComputeBoundary(const std::vector &objects) {
  	//初始化值为正无穷大
    min_x_ = std::numeric_limits::infinity();
    min_y_ = std::numeric_limits::infinity();
    max_x_ = -std::numeric_limits::infinity();
    max_y_ = -std::numeric_limits::infinity();
    for (ObjectPtr object : objects) {
      min_x_ = std::fmin(min_x_, object->aabox().min_x());
      max_x_ = std::fmax(max_x_, object->aabox().max_x());
      min_y_ = std::fmin(min_y_, object->aabox().min_y());
      max_y_ = std::fmax(max_y_, object->aabox().max_y());
    }
    mid_x_ = (min_x_ + max_x_) / 2.0;
    mid_y_ = (min_y_ + max_y_) / 2.0;
    CHECK(!std::isinf(max_x_) && !std::isinf(max_y_) && !std::isinf(min_x_) &&
          !std::isinf(min_y_))
        << "the provided object box size is infinity";
  }

step 2 计算按哪个方向划分以及划分的点（作为当前节点）

KD-Tree划分计算：长和宽哪个更大就按照哪个方向进行划分，划分的点也比较粗暴，直接采用中点进行划分，因此这里和笔记中的选择划分点有区别，笔记中或一般 kd-tree 的节点（划分点）选取依然是输入数据中的点，而这里不是，这里直接以物理空间进行“均分”。

  void ComputePartition() {
    if (max_x_ - min_x_ >= max_y_ - min_y_) {
      partition_ = PARTITION_X;
      partition_position_ = (min_x_ + max_x_) / 2.0;
    } else {
      partition_ = PARTITION_Y;
      partition_position_ = (min_y_ + max_y_) / 2.0;
    }
  }

step 3 分裂当前空间（Node）为左右子空间（左右子树）

在 step 3.1中构建了当前节点的数据，于是看可以看出建树的过程是一个前序遍历建树的过程。

    if (SplitToSubNodes(objects, params)) {//检查是否分裂是否满足参数条件
      std::vector left_subnode_objects;
      std::vector right_subnode_objects;
      //Step 3.1.划分空间
      PartitionObjects(objects, &left_subnode_objects, &right_subnode_objects);

      //Step 3.2. Split to sub-nodes. 分裂到左右子树
      if (!left_subnode_objects.empty()) {
      	// 在构造函数中递归
        left_subnode_.reset(new AABoxKDTree2dNode(
            left_subnode_objects, params, depth + 1));
      }
      if (!right_subnode_objects.empty()) {
        right_subnode_.reset(new AABoxKDTree2dNode(
            right_subnode_objects, params, depth + 1));
      }
    }

首先检测是否需要继续分裂：

  bool SplitToSubNodes(const std::vector &objects,
                       const AABoxKDTreeParams ¶ms) {
    if (params.max_depth >= 0 && depth_ >= params.max_depth) {
      return false;
    }
    if (static_cast(objects.size()) <= std::max(1, params.max_leaf_size)) {
      return false;
    }
    // 这里的维度距离就可以看作是像素之间的距离
    if (params.max_leaf_dimension >= 0.0 &&
        std::max(max_x_ - min_x_, max_y_ - min_y_) <=
            params.max_leaf_dimension) {
      return false;
    }
    return true;
  }

开始划分左右子空间：

  //将当前Node分裂为左右子树
  void PartitionObjects(const std::vector &objects,
                        std::vector *const left_subnode_objects,
                        std::vector *const right_subnode_objects) {
    left_subnode_objects->clear();
    right_subnode_objects->clear();
    std::vector other_objects;
    if (partition_ == PARTITION_X) {
      for (ObjectPtr object : objects) {
        if (object->aabox().max_x() <= partition_position_) {
          left_subnode_objects->push_back(object);
        } else if (object->aabox().min_x() >= partition_position_) {
          right_subnode_objects->push_back(object);
        } else {
          other_objects.push_back(object);//在分界线上
        }
      }
    } else {
      for (ObjectPtr object : objects) {
        if (object->aabox().max_y() <= partition_position_) {
          left_subnode_objects->push_back(object);
        } else if (object->aabox().min_y() >= partition_position_) {
          right_subnode_objects->push_back(object);
        } else {
          other_objects.push_back(object);
        }
      }
    }
    // 保存当前节点数据
    InitObjects(other_objects);
  }

  //$ 对当前节点数据进行处理保存
  void InitObjects(const std::vector &objects) {
    num_objects_ = static_cast(objects.size());  //当前节点个数？
    objects_sorted_by_min_ = objects;
    objects_sorted_by_max_ = objects;
    std::sort(objects_sorted_by_min_.begin(), objects_sorted_by_min_.end(),
              [&](ObjectPtr obj1, ObjectPtr obj2) {
                return partition_ == PARTITION_X
                           ? obj1->aabox().min_x() < obj2->aabox().min_x()
                           : obj1->aabox().min_y() < obj2->aabox().min_y();
              });
    std::sort(objects_sorted_by_max_.begin(), objects_sorted_by_max_.end(),
              [&](ObjectPtr obj1, ObjectPtr obj2) {
                return partition_ == PARTITION_X
                           ? obj1->aabox().max_x() > obj2->aabox().max_x()
                           : obj1->aabox().max_y() > obj2->aabox().max_y();
              });
    objects_sorted_by_min_bound_.reserve(
        num_objects_);  //记录从小到大排序的坐标x值
    for (ObjectPtr object : objects_sorted_by_min_) {
      objects_sorted_by_min_bound_.push_back(partition_ == PARTITION_X
                                                 ? object->aabox().min_x()
                                                 : object->aabox().min_y());
    }
    objects_sorted_by_max_bound_.reserve(num_objects_);
    for (ObjectPtr object : objects_sorted_by_max_) {
      objects_sorted_by_max_bound_.push_back(partition_ == PARTITION_X
                                                 ? object->aabox().max_x()
                                                 : object->aabox().max_y());
    }
  }

2.查询

构建树的目的就是降低查找的时间复杂度。可以查找最近邻，也可以查找k近邻。

（1）最近邻查找

1.首先要找到该目标点的叶子节点进行判断，记录最近距离与nearest_object;

2.回溯判断另一半空间是否需要计算，若需要在另一个子树继续搜索最近邻，更新记录值。

3.当回溯到根节点时，算法结束，此时保存的最近邻节点就是最终的最近邻。

参考阅读详细看查询距离。

void GetNearestObjectInternal(
      const Vec2d &point,
      double *const min_distance_sqr,  //指针不变，但指针指向的内容可以变
      ObjectPtr *const nearest_object) const {
    // Step 4.距离大于当前最小距离的子树就不计算了，直接返回（回溯）-->剪枝操作
    if (LowerDistanceSquareToPoint(point) >= *min_distance_sqr - kMathEpsilon) {
      return;
    }
    // !后序访问：注意这里不是遍历访问
    // Step 1.DFS:将按point搜索，改为按当前节点划分方式的值进行搜索
    const double pvalue = (partition_ == PARTITION_X ? point.x() : point.y());
    const bool search_left_first = (pvalue < partition_position_);
    // 这里的 DFS 不是遍历访问，是 if else 关系，访问另一子空间是在回溯时访问
    if (search_left_first) {
      if (left_subnode_ != nullptr) {
        left_subnode_->GetNearestObjectInternal(point, min_distance_sqr,
                                                nearest_object);  //递归查找
      }
    } else {
      if (right_subnode_ != nullptr) {
        right_subnode_->GetNearestObjectInternal(point, min_distance_sqr,
                                                 nearest_object);
      }
    }
    // 距离 == 0 ， 找到最近邻，直接返回，不用深入了
    if (*min_distance_sqr <= kMathEpsilon) {  //回溯的时候min_distance_sqr值会改变
      return;
    }

    // Step 2.欧氏距离找最近的object
    if (search_left_first) {
      for (int i = 0; i < num_objects_; ++i) {
        const double bound = objects_sorted_by_min_bound_[i];
        // ?框的最小边界都比它大，说明在框外面的意思？
        if (bound > pvalue && Square(bound - pvalue) > *min_distance_sqr) {
          break;
        }
        ObjectPtr object = objects_sorted_by_min_[i];
        const double distance_sqr = object->DistanceSquareTo(point);
        if (distance_sqr < *min_distance_sqr) {
          *min_distance_sqr = distance_sqr;
          *nearest_object = object;
        }
      }
    } else {
      for (int i = 0; i < num_objects_; ++i) {
        const double bound = objects_sorted_by_max_bound_[i];
        // ?框的最大边界都比它小，说明在框外面的意思？
        if (bound < pvalue && Square(bound - pvalue) > *min_distance_sqr) {
          break;
        }
        ObjectPtr object = objects_sorted_by_max_[i];
        const double distance_sqr = object->DistanceSquareTo(point);
        if (distance_sqr < *min_distance_sqr) {
          *min_distance_sqr = distance_sqr;
          *nearest_object = object;
        }
      }
    }

    // 距离 == 0 ， 找到最近邻，直接返回，不用深入了
    if (*min_distance_sqr <= kMathEpsilon) {
      return;
    }

    // Step 3.回溯时访问另一半子空间，看看是否还有距离更近的点
    if (search_left_first) {
      if (right_subnode_ != nullptr) {
        right_subnode_->GetNearestObjectInternal(point, min_distance_sqr,
                                                 nearest_object);
      }
    } else {
      if (left_subnode_ != nullptr) {
        left_subnode_->GetNearestObjectInternal(point, min_distance_sqr,
                                                nearest_object);
      }
    }
  }

按距离查找原理类似。

Apollo | common | math | kd-tree

C/C++/C#相关栏目本月热门文章