FindBestSplit of C++ version may not find the best split

[Datou0718]

Hi:

I just traced through the code and found that the splitting algorithm may not yield the best result.
This is the original code in TreeClassification::findBestSplitValueSmallQ:

        // Sum of squares
        double sum_left = 0;
        double sum_right = 0;
        for (size_t j = 0; j < num_classes; ++j)
        {
          class_counts_left[j] += counter_per_class[i * num_classes + j];
          size_t class_count_right = class_counts[j] - class_counts_left[j];

          sum_left += (*class_weights)[j] * class_counts_left[j] * class_counts_left[j];
          sum_right += (*class_weights)[j] * class_count_right * class_count_right;

        }
        // Decrease of impurity
        decrease = sum_right / (double)n_right + sum_left / (double)n_left;

You can see that the class counts are squared, while only divided by n_right (or n_left) once. This is kind of like the Gini Impurity but not exactly the same. I’m not sure but I think this may lead to a suboptimal split. The problem can also be seen in other findBestSplit functions as well as addGiniImportance.
Below is my revised version for your reference:

        // Sum of squares
        double sum_left = 0;
        double sum_right = 0;
        for (size_t j = 0; j < num_classes; ++j)
        {
          class_counts_left[j] += counter_per_class[i * num_classes + j];
          size_t class_count_right = class_counts[j] - class_counts_left[j];

          sum_left += class_counts_left[j] * class_counts_left[j] / (double)(n_left * n_left);
          sum_right += class_count_right * class_count_right / (double)(n_right * n_right);
        }
        double original_impurity = 1;
        for (size_t j = 0; j < num_classes; ++j)
        {
          original_impurity -= class_counts[j] * class_counts[j] / (double)(num_samples_node * num_samples_node);
        }
        decrease = original_impurity - (1.0 - sum_left * (double)n_left / (double)num_samples_node - sum_right * (double)n_right / (double)num_samples_node);

[mnwright]

Thanks. I'm pretty sure that the calculation is correct. Note that some terms cancel out and we don't need all terms for optimization (hence we calculate those only for the impurity importance).