diff options
Diffstat (limited to 'geom_matching/wasserstein/include/dnn/local')
-rw-r--r-- | geom_matching/wasserstein/include/dnn/local/kd-tree.h | 2 | ||||
-rw-r--r-- | geom_matching/wasserstein/include/dnn/local/kd-tree.hpp | 29 |
2 files changed, 23 insertions, 8 deletions
diff --git a/geom_matching/wasserstein/include/dnn/local/kd-tree.h b/geom_matching/wasserstein/include/dnn/local/kd-tree.h index 7e01072..13eaf27 100644 --- a/geom_matching/wasserstein/include/dnn/local/kd-tree.h +++ b/geom_matching/wasserstein/include/dnn/local/kd-tree.h @@ -49,7 +49,7 @@ namespace dnn void init(const Range& range); DistanceType weight(PointHandle p) { return weights_[indices_[p]]; } - void increase_weight(PointHandle p, DistanceType w); + void change_weight(PointHandle p, DistanceType w); HandleDistance find(PointHandle q) const; Result findR(PointHandle q, DistanceType r) const; // all neighbors within r diff --git a/geom_matching/wasserstein/include/dnn/local/kd-tree.hpp b/geom_matching/wasserstein/include/dnn/local/kd-tree.hpp index 6b0852c..22108aa 100644 --- a/geom_matching/wasserstein/include/dnn/local/kd-tree.hpp +++ b/geom_matching/wasserstein/include/dnn/local/kd-tree.hpp @@ -164,11 +164,15 @@ search(PointHandle q, ResultsFunctor& rf) const template<class T> void dnn::KDTree<T>:: -increase_weight(PointHandle p, DistanceType w) +change_weight(PointHandle p, DistanceType w) { size_t idx = indices_[p]; - // weight should only increase - assert( weights_[idx] <= w ); + + if ( weights_[idx] == w ) { + return; + } + + bool weight_increases = ( weights_[idx] < w ); weights_[idx] = w; typedef std::tuple<HCIterator, HCIterator> KDTreeNode; @@ -223,10 +227,21 @@ increase_weight(PointHandle p, DistanceType w) min_w = weights_[im]; } - if (subtree_weights_[im] < min_w ) // increase weight - subtree_weights_[im] = min_w; - else - break; + if (weight_increases) { + + if (subtree_weights_[im] < min_w ) // increase weight + subtree_weights_[im] = min_w; + else + break; + + } else { + + if (subtree_weights_[im] > min_w ) // decrease weight + subtree_weights_[im] = min_w; + else + break; + + } } } |