gini_gain.hpp

Go to the documentation of this file.

#ifndef MLPACK_METHODS_DECISION_TREE_GINI_GAIN_HPP
#define MLPACK_METHODS_DECISION_TREE_GINI_GAIN_HPP

#include <mlpack/core.hpp>

namespace mlpack {
namespace tree {

class GiniGain
{
 public:
  template<typename RowType>
  static double Evaluate(const RowType& labels,
                         const size_t numClasses)
  {
    // Corner case: if there are no elements, the impurity is zero.
    if (labels.n_elem == 0)
      return 0.0;

    arma::Col<size_t> counts(numClasses);
    counts.zeros();
    for (size_t i = 0; i < labels.n_elem; ++i)
      counts[labels[i]]++;

    // Calculate the Gini impurity of the un-split node.
    double impurity = 0.0;
    for (size_t i = 0; i < numClasses; ++i)
    {
      const double f = ((double) counts[i] / (double) labels.n_elem);
      impurity += f * (1.0 - f);
    }

    return -impurity;
  }

  static double Range(const size_t numClasses)
  {
    // The best possible case is that only one class exists, which gives a Gini
    // impurity of 0.  The worst possible case is that the classes are evenly
    // distributed, which gives n * (1/n * (1 - 1/n)) = 1 - 1/n.
    return 1.0 - (1.0 / double(numClasses));
  }
};

} // namespace tree
} // namespace mlpack

#endif