[if-else] use bit operation for categorical features (#821)

* bitset for if-else * add const
microsoft · Aug 16, 2017 · f8d6db9 · f8d6db9
1 parent a77920c
commit f8d6db9
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 17 deletions.
diff --git a/include/LightGBM/tree.h b/include/LightGBM/tree.h
@@ -238,24 +238,11 @@ class Tree {
     } else {
       str_buf << "if (std::isnan(fval)) { int_fval = 0; } else { int_fval = static_cast<int>(fval); }";
     }
-    str_buf << "if (";
     int cat_idx = int(threshold_[node]);
-    // To-do: optimize this by using bitset
-    std::vector<int> cats;
-    for (int i = cat_boundaries_[cat_idx]; i < cat_boundaries_[cat_idx + 1]; ++i) {
-      for (int j = 0; j < 32; ++j) {
-        int cat = (i - cat_boundaries_[cat_idx]) * 32 + j;
-        if (Common::FindInBitset(cat_threshold_.data() + cat_boundaries_[cat_idx],
-                                 cat_boundaries_[cat_idx + 1] - cat_boundaries_[cat_idx], cat)) {
-          cats.push_back(cat);
-        }
-      }
-    }
-    str_buf << "int_fval >= 0 && (";
-    for (int i = 0; i < static_cast<int>(cats.size()) - 1; ++i) {
-      str_buf << "int_fval == " << cats[i] << " || ";
-    }
-    str_buf << "int_fval == " << cats.back() << ")) {";
+    str_buf << "if (int_fval >= 0 && int_fval < 32 * (";
+    str_buf << cat_boundaries_[cat_idx + 1] - cat_boundaries_[cat_idx];
+    str_buf << ") && (((cat_threshold[" << cat_boundaries_[cat_idx];
+    str_buf << " + int_fval / 32] >> (int_fval & 31)) & 1))) {";
     return str_buf.str();
   }
 

diff --git a/src/io/tree.cpp b/src/io/tree.cpp
@@ -310,6 +310,14 @@ std::string Tree::ToIfElse(int index, bool is_predict_leaf_index) {
   if (num_leaves_ == 1) {
     str_buf << "return 0";
   } else {
+    str_buf << "const std::vector<uint32_t> cat_threshold = {";
+    for (size_t i = 0; i < cat_threshold_.size(); ++i) {
+      if (i != 0) {
+        str_buf << ",";
+      }
+      str_buf << cat_threshold_[i];
+    }
+    str_buf << "};";
     // use this for the missing value conversion
     str_buf << "double fval = 0.0f; ";
     if (num_cat_ > 0) {