fix extratrees splitting with num.random.splits>1

adamwanggang · Aug 22, 2019 · ce29a66 · ce29a66
1 parent d15cf25
commit ce29a66
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 0 deletions.
diff --git a/src/TreeClassification.cpp b/src/TreeClassification.cpp
@@ -492,6 +492,9 @@ void TreeClassification::findBestSplitValueExtraTrees(size_t nodeID, size_t varI
   for (size_t i = 0; i < num_random_splits; ++i) {
     possible_split_values.push_back(udist(random_number_generator));
   }
+  if (num_random_splits > 1) {
+    std::sort(possible_split_values.begin(), possible_split_values.end());
+  }
 
   const size_t num_splits = possible_split_values.size();
   if (memory_saving_splitting) {

diff --git a/src/TreeProbability.cpp b/src/TreeProbability.cpp
@@ -496,6 +496,9 @@ void TreeProbability::findBestSplitValueExtraTrees(size_t nodeID, size_t varID,
   for (size_t i = 0; i < num_random_splits; ++i) {
     possible_split_values.push_back(udist(random_number_generator));
   }
+  if (num_random_splits > 1) {
+    std::sort(possible_split_values.begin(), possible_split_values.end());
+  }
 
   const size_t num_splits = possible_split_values.size();
   if (memory_saving_splitting) {

diff --git a/src/TreeRegression.cpp b/src/TreeRegression.cpp
@@ -541,6 +541,9 @@ void TreeRegression::findBestSplitValueExtraTrees(size_t nodeID, size_t varID, d
   for (size_t i = 0; i < num_random_splits; ++i) {
     possible_split_values.push_back(udist(random_number_generator));
   }
+  if (num_random_splits > 1) {
+    std::sort(possible_split_values.begin(), possible_split_values.end());
+  }
 
   const size_t num_splits = possible_split_values.size();
   if (memory_saving_splitting) {

diff --git a/src/TreeSurvival.cpp b/src/TreeSurvival.cpp
@@ -708,6 +708,9 @@ void TreeSurvival::findBestSplitValueExtraTrees(size_t nodeID, size_t varID, dou
   for (size_t i = 0; i < num_random_splits; ++i) {
     possible_split_values.push_back(udist(random_number_generator));
   }
+  if (num_random_splits > 1) {
+    std::sort(possible_split_values.begin(), possible_split_values.end());
+  }
 
   size_t num_splits = possible_split_values.size();