As some processed distances are already ^2, use template to select whether or not we have to ^2 in KMeanspp

Author: pemmanuelviel

src/cpp/flann/algorithms/center_chooser.h

@@ -13,6 +13,63 @@
 namespace flann
 {
 
+template <typename Distance, typename ElementType>
+struct squareDistance
+{
+    typedef typename Distance::ResultType ResultType;
+    ResultType operator()( ResultType dist ) { return dist*dist; }
+};
+
+
+template <typename ElementType>
+struct squareDistance<L2_Simple<ElementType>, ElementType>
+{
+    typedef typename L2_Simple<ElementType>::ResultType ResultType;
+    ResultType operator()( ResultType dist ) { return dist; }
+};
+
+template <typename ElementType>
+struct squareDistance<L2_3D<ElementType>, ElementType>
+{
+    typedef typename L2_3D<ElementType>::ResultType ResultType;
+    ResultType operator()( ResultType dist ) { return dist; }
+};
+
+template <typename ElementType>
+struct squareDistance<L2<ElementType>, ElementType>
+{
+    typedef typename L2<ElementType>::ResultType ResultType;
+    ResultType operator()( ResultType dist ) { return dist; }
+};
+
+
+template <typename ElementType>
+struct squareDistance<HellingerDistance<ElementType>, ElementType>
+{
+    typedef typename HellingerDistance<ElementType>::ResultType ResultType;
+    ResultType operator()( ResultType dist ) { return dist; }
+};
+
+
+template <typename ElementType>
+struct squareDistance<ChiSquareDistance<ElementType>, ElementType>
+{
+    typedef typename ChiSquareDistance<ElementType>::ResultType ResultType;
+    ResultType operator()( ResultType dist ) { return dist; }
+};
+
+
+template <typename Distance>
+typename Distance::ResultType ensureSquareDistance( typename Distance::ResultType dist )
+{
+    typedef typename Distance::ElementType ElementType;
+
+    squareDistance<Distance, ElementType> dummy;
+    return dummy( dist );
+}
+
+
+
 template <typename Distance>
 class CenterChooser
 {
@@ -180,7 +237,7 @@ class KMeansppCenterChooser : public CenterChooser<Distance>
         // far from previous centers (and this complies to "k-means++: the advantages of careful seeding" article)
         for (int i = 0; i < n; i++) {
             closestDistSq[i] = distance_(points_[indices[i]], points_[indices[index]], cols_);
-            closestDistSq[i] *= closestDistSq[i];
+            closestDistSq[i] = ensureSquareDistance<Distance>( closestDistSq[i] );
             currentPot += closestDistSq[i];
         }
 
@@ -208,7 +265,7 @@ class KMeansppCenterChooser : public CenterChooser<Distance>
                 double newPot = 0;
                 for (int i = 0; i < n; i++) {
                     DistanceType dist = distance_(points_[indices[i]], points_[indices[index]], cols_);
-                    newPot += std::min( dist*dist, closestDistSq[i] );
+                    newPot += std::min( ensureSquareDistance<Distance>(dist), closestDistSq[i] );
                 }
 
                 // Store the best result
@@ -223,7 +280,7 @@ class KMeansppCenterChooser : public CenterChooser<Distance>
             currentPot = bestNewPot;
             for (int i = 0; i < n; i++) {
                 DistanceType dist = distance_(points_[indices[i]], points_[indices[bestNewIndex]], cols_);
-                closestDistSq[i] = std::min( dist*dist, closestDistSq[i] );
+                closestDistSq[i] = std::min( ensureSquareDistance<Distance>(dist), closestDistSq[i] );
             }
         }