diff --git a/src/mlpack/methods/neighbor_search/neighbor_search_rules_impl.hpp b/src/mlpack/methods/neighbor_search/neighbor_search_rules_impl.hpp index 031965c36c..639483b996 100644 --- a/src/mlpack/methods/neighbor_search/neighbor_search_rules_impl.hpp +++ b/src/mlpack/methods/neighbor_search/neighbor_search_rules_impl.hpp @@ -183,7 +183,11 @@ inline double NeighborSearchRules::Score( } else { - adjustedScore = SortPolicy::CombineBest(adjustedScore, queryDescDist); + // If the parent node is NULL, force adjustedScore to be such that it can't + // be pruned. Otherwise use the parent descendant distance. + const double queryParentDescDist = (queryNode.Parent() == NULL) ? + bestDistance : queryNode.Parent()->FurthestDescendantDistance(); + adjustedScore = SortPolicy::CombineBest(adjustedScore, queryParentDescDist); } if (traversalInfo.LastReferenceNode() == referenceNode.Parent()) @@ -193,7 +197,11 @@ inline double NeighborSearchRules::Score( } else { - adjustedScore = SortPolicy::CombineBest(adjustedScore, refDescDist); + // If the parent node is NULL, force adjustedScore to be such that it can't + // be pruned. Otherwise use the parent descendant distance. + const double refParentDescDist = (referenceNode.Parent() == NULL) ? + bestDistance : referenceNode.Parent()->FurthestDescendantDistance(); + adjustedScore = SortPolicy::CombineBest(adjustedScore, refParentDescDist); } // Can we prune?