From 670a56a58f050eb2a34a91028198b4e282a4fc6e Mon Sep 17 00:00:00 2001
From: Ryan Curtin <ryan@ratml.org>
Date: Tue, 17 Feb 2015 14:52:32 -0500
Subject: [PATCH] Refactor to apply mappings earlier.

---
 src/mlpack/methods/kmeans/dtnn_kmeans_impl.hpp | 13 +++----------
 src/mlpack/methods/kmeans/dtnn_rules_impl.hpp  |  7 +++++--
 2 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/src/mlpack/methods/kmeans/dtnn_kmeans_impl.hpp b/src/mlpack/methods/kmeans/dtnn_kmeans_impl.hpp
index dd6665b696..08482f8177 100644
--- a/src/mlpack/methods/kmeans/dtnn_kmeans_impl.hpp
+++ b/src/mlpack/methods/kmeans/dtnn_kmeans_impl.hpp
@@ -181,9 +181,7 @@ void DTNNKMeans<MetricType, MatType, TreeType>::UpdateTree(
   if ((node.Stat().Pruned() == clusterDistances.n_elem - 1) &&
       (node.Stat().Owner() < clusterDistances.n_elem - 1))
   {
-    const size_t owner = (tree::TreeTraits<TreeType>::RearrangesDataset) ?
-        oldFromNewCentroids[node.Stat().Owner()] :
-        node.Stat().Owner();
+    const size_t owner = node.Stat().Owner();
 
     node.Stat().LastUpperBound() = node.Stat().UpperBound() +
         clusterDistances[owner];
@@ -211,7 +209,6 @@ void DTNNKMeans<MetricType, MatType, TreeType>::UpdateTree(
   }
   node.Stat().Pruned() = size_t(-1);
   node.Stat().Owner() = size_t(-1);
-  node.Stat().LowerBound() = DBL_MAX;
 
   for (size_t i = 0; i < node.NumChildren(); ++i)
     UpdateTree(node.Child(i), clusterDistances, oldFromNewCentroids);
@@ -227,9 +224,7 @@ void DTNNKMeans<MetricType, MatType, TreeType>::ExtractCentroids(
   // Does this node own points?
   if (node.Stat().Pruned() == newCentroids.n_cols)
   {
-    const size_t owner = (tree::TreeTraits<TreeType>::RearrangesDataset) ?
-        oldFromNewCentroids[node.Stat().Owner()] :
-        node.Stat().Owner();
+    const size_t owner = node.Stat().Owner();
     newCentroids.col(owner) += node.Stat().Centroid() * node.NumDescendants();
     newCounts[owner] += node.NumDescendants();
   }
@@ -241,9 +236,7 @@ void DTNNKMeans<MetricType, MatType, TreeType>::ExtractCentroids(
     {
       for (size_t i = 0; i < node.NumPoints(); ++i)
       {
-        const size_t owner = (tree::TreeTraits<TreeType>::RearrangesDataset) ?
-            oldFromNewCentroids[assignments[node.Point(i)]] :
-            assignments[node.Point(i)];
+        const size_t owner = assignments[node.Point(i)];
         newCentroids.col(owner) += dataset.col(node.Point(i));
         ++newCounts[owner];
       }
diff --git a/src/mlpack/methods/kmeans/dtnn_rules_impl.hpp b/src/mlpack/methods/kmeans/dtnn_rules_impl.hpp
index 6a57273a39..d5268f4e36 100644
--- a/src/mlpack/methods/kmeans/dtnn_rules_impl.hpp
+++ b/src/mlpack/methods/kmeans/dtnn_rules_impl.hpp
@@ -58,7 +58,8 @@ inline force_inline double DTNNKMeansRules<MetricType, TreeType>::BaseCase(
   {
     lowerBounds[queryIndex] = upperBounds[queryIndex];
     upperBounds[queryIndex] = distance;
-    assignments[queryIndex] = referenceIndex;
+    assignments[queryIndex] = (tree::TreeTraits<TreeType>::RearrangesDataset) ?
+        oldFromNewCentroids[referenceIndex] : referenceIndex;
   }
   else if (distance < lowerBounds[queryIndex])
   {
@@ -121,7 +122,9 @@ inline double DTNNKMeansRules<MetricType, TreeType>::Score(
     queryNode.Stat().UpperBound() = distances.Hi();
     // If this node has only one descendant, then it may be the owner.
     if (referenceNode.NumDescendants() == 1)
-      queryNode.Stat().Owner() = referenceNode.Descendant(0);
+      queryNode.Stat().Owner() = (tree::TreeTraits<TreeType>::RearrangesDataset)
+          ? oldFromNewCentroids[referenceNode.Descendant(0)]
+          : referenceNode.Descendant(0);
   }
 
   // Is everything pruned?