Merge pull request #641 from MarcosPividori/master

Properly use Enum type.
mlpack · May 23, 2016 · 9b42c22 · 9b42c22
2 parents 39eefde + dd136db
commit 9b42c22
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 7 deletions.
diff --git a/src/mlpack/methods/neighbor_search/kfn_main.cpp b/src/mlpack/methods/neighbor_search/kfn_main.cpp
@@ -149,7 +149,7 @@ int main(int argc, char *argv[])
     const string treeType = CLI::GetParam<string>("tree_type");
     const bool randomBasis = CLI::HasParam("random_basis");
 
-    int tree = 0;
+    KFNModel::TreeTypes tree = KFNModel::KD_TREE;
     if (treeType == "kd")
       tree = KFNModel::KD_TREE;
     else if (treeType == "cover")

diff --git a/src/mlpack/methods/neighbor_search/knn_main.cpp b/src/mlpack/methods/neighbor_search/knn_main.cpp
@@ -153,7 +153,7 @@ int main(int argc, char *argv[])
     const string treeType = CLI::GetParam<string>("tree_type");
     const bool randomBasis = CLI::HasParam("random_basis");
 
-    int tree = 0;
+    KNNModel::TreeTypes tree = KNNModel::KD_TREE;
     if (treeType == "kd")
       tree = KNNModel::KD_TREE;
     else if (treeType == "cover")

diff --git a/src/mlpack/methods/neighbor_search/ns_model.hpp b/src/mlpack/methods/neighbor_search/ns_model.hpp
@@ -52,7 +52,7 @@ class NSModel
   };
 
  private:
-  int treeType;
+  TreeTypes treeType;
   size_t leafSize;
 
   // For random projections.
@@ -83,7 +83,7 @@ class NSModel
    * Initialize the NSModel with the given type and whether or not a random
    * basis should be used.
    */
-  NSModel(int treeType = TreeTypes::KD_TREE, bool randomBasis = false);
+  NSModel(TreeTypes treeType = TreeTypes::KD_TREE, bool randomBasis = false);
 
   //! Clean memory, if necessary.
   ~NSModel();
@@ -105,8 +105,8 @@ class NSModel
   size_t LeafSize() const { return leafSize; }
   size_t& LeafSize() { return leafSize; }
 
-  int TreeType() const { return treeType; }
-  int& TreeType() { return treeType; }
+  TreeTypes TreeType() const { return treeType; }
+  TreeTypes& TreeType() { return treeType; }
 
   bool RandomBasis() const { return randomBasis; }
   bool& RandomBasis() { return randomBasis; }

diff --git a/src/mlpack/methods/neighbor_search/ns_model_impl.hpp b/src/mlpack/methods/neighbor_search/ns_model_impl.hpp
@@ -21,7 +21,7 @@ namespace neighbor {
  * basis should be used.
  */
 template<typename SortPolicy>
-NSModel<SortPolicy>::NSModel(int treeType, bool randomBasis) :
+NSModel<SortPolicy>::NSModel(TreeTypes treeType, bool randomBasis) :
     treeType(treeType),
     randomBasis(randomBasis),
     kdTreeNS(NULL),