[mlpack-git] master: Do mapping correctly; handle cover trees right. (2201a2f)

[gitdub at big.cc.gt.atl.ga.us]

Repository : https://github.com/mlpack/mlpack

On branch  : master
Link       : https://github.com/mlpack/mlpack/compare/eddd7167d69b6c88b271ef2e51d1c20e13f1acd8...70342dd8e5c17e0c164cfb8189748671e9c0dd44

>---------------------------------------------------------------

commit 2201a2f05d0cb2abe915be3dff021b5ab81f7d00
Author: Ryan Curtin <ryan at ratml.org>
Date:   Tue Feb 17 14:22:16 2015 -0500

    Do mapping correctly; handle cover trees right.


>---------------------------------------------------------------

2201a2f05d0cb2abe915be3dff021b5ab81f7d00
 src/mlpack/methods/kmeans/dtnn_kmeans_impl.hpp | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/src/mlpack/methods/kmeans/dtnn_kmeans_impl.hpp b/src/mlpack/methods/kmeans/dtnn_kmeans_impl.hpp
index 38e7650..dd6665b 100644
--- a/src/mlpack/methods/kmeans/dtnn_kmeans_impl.hpp
+++ b/src/mlpack/methods/kmeans/dtnn_kmeans_impl.hpp
@@ -181,7 +181,9 @@ 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 = oldFromNewCentroids[node.Stat().Owner()];
+    const size_t owner = (tree::TreeTraits<TreeType>::RearrangesDataset) ?
+        oldFromNewCentroids[node.Stat().Owner()] :
+        node.Stat().Owner();
 
     node.Stat().LastUpperBound() = node.Stat().UpperBound() +
         clusterDistances[owner];
@@ -225,18 +227,26 @@ void DTNNKMeans<MetricType, MatType, TreeType>::ExtractCentroids(
   // Does this node own points?
   if (node.Stat().Pruned() == newCentroids.n_cols)
   {
-    const size_t owner = oldFromNewCentroids[node.Stat().Owner()];
+    const size_t owner = (tree::TreeTraits<TreeType>::RearrangesDataset) ?
+        oldFromNewCentroids[node.Stat().Owner()] :
+        node.Stat().Owner();
     newCentroids.col(owner) += node.Stat().Centroid() * node.NumDescendants();
     newCounts[owner] += node.NumDescendants();
   }
   else
   {
     // Check each point held in the node.
-    for (size_t i = 0; i < node.NumPoints(); ++i)
+    // Only check at leaves.
+    if (node.NumChildren() == 0)
     {
-      const size_t owner = oldFromNewCentroids[assignments[node.Point(i)]];
-      newCentroids.col(owner) += dataset.col(node.Point(i));
-      ++newCounts[owner];
+      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)];
+        newCentroids.col(owner) += dataset.col(node.Point(i));
+        ++newCounts[owner];
+      }
     }
 
     // The node is not entirely owned by a cluster.  Recurse.