From 030d36d7fa6fd9a5d6fb317ed244c5de9a8b866c Mon Sep 17 00:00:00 2001
From: Ryan Curtin <ryan@ratml.org>
Date: Thu, 19 Feb 2015 14:21:40 -0500
Subject: [PATCH] Implement Rescore(). Minor speedup.

---
 src/mlpack/methods/kmeans/dtnn_rules_impl.hpp | 34 +++++++++++++++----
 1 file changed, 28 insertions(+), 6 deletions(-)

diff --git a/src/mlpack/methods/kmeans/dtnn_rules_impl.hpp b/src/mlpack/methods/kmeans/dtnn_rules_impl.hpp
index 195cd1a4d5..ca8943c432 100644
--- a/src/mlpack/methods/kmeans/dtnn_rules_impl.hpp
+++ b/src/mlpack/methods/kmeans/dtnn_rules_impl.hpp
@@ -78,7 +78,8 @@ inline double DTNNKMeansRules<MetricType, TreeType>::Score(
   if (prunedPoints[queryIndex])
     return DBL_MAX;
 
-  // No pruning at this level (for now).
+  // No pruning at this level; we're not likely to encounter a single query
+  // point with a reference node..
   return 0;
 }
 
@@ -99,9 +100,7 @@ inline double DTNNKMeansRules<MetricType, TreeType>::Score(
   }
 
   if (queryNode.Stat().Pruned() == centroids.n_cols)
-  {
     return DBL_MAX;
-  }
 
   // Get minimum and maximum distances.
   math::Range distances = queryNode.RangeDistance(&referenceNode);
@@ -151,11 +150,34 @@ inline double DTNNKMeansRules<MetricType, TreeType>::Rescore(
 
 template<typename MetricType, typename TreeType>
 inline double DTNNKMeansRules<MetricType, TreeType>::Rescore(
-    TreeType& /* queryNode */,
-    TreeType& /* referenceNode */,
+    TreeType& queryNode,
+    TreeType& referenceNode,
     const double oldScore)
 {
-  // No rescoring (for now).
+  if (oldScore == DBL_MAX)
+    return DBL_MAX; // It's already pruned.
+
+  // oldScore contains the minimum distance between queryNode and referenceNode.
+  // In the time since Score() has been called, the upper bound *may* have
+  // tightened.  If it has tightened enough, we may prune this node now.
+  if (oldScore > queryNode.Stat().UpperBound())
+  {
+    // We may still be able to improve the lower bound on pruned nodes.
+    if (oldScore < queryNode.Stat().LowerBound())
+      queryNode.Stat().LowerBound() = oldScore;
+
+    // This assumes that reference clusters don't appear elsewhere in the tree.
+    queryNode.Stat().Pruned() += referenceNode.NumDescendants();
+    return DBL_MAX;
+  }
+
+  // Also, check if everything has been pruned.
+  if (queryNode.Stat().Pruned() == centroids.n_cols - 1)
+  {
+    queryNode.Stat().Pruned() = centroids.n_cols; // Owner() is already set.
+    return DBL_MAX;
+  }
+
   return oldScore;
 }