Criterion.cpp (880B)
1 #include "Criterion.h" 2 #include <unordered_map> 3 #include <math.h> 4 5 float Criterion::giniImpurity(float* X, int* y, int samples, int features, int index, float splitValue){ 6 7 std::unordered_map<int, int> ltMap; 8 std::unordered_map<int, int> gtMap; 9 10 int ltCount = 0; 11 int gteqCount = 0; 12 13 for(int i = 0; i < samples; ++i){ 14 if(X[index + (i * features)] < splitValue){ 15 ltMap[y[i]]++; 16 ltCount++; 17 } 18 else{ 19 gtMap[y[i]]++; 20 gteqCount++; 21 } 22 } 23 24 25 float ltGini= 1.0f; 26 27 for (const auto& pair : ltMap) { 28 ltGini -= pow(float(pair.second) / ltCount, 2); 29 } 30 31 float gteqGini = 1.0f; 32 33 for (const auto& pair : gtMap) { 34 gteqGini -= pow(float(pair.second) / gteqCount, 2); 35 } 36 37 if(gteqCount == 0){ 38 gteqGini = 0.0f; 39 } 40 if(ltCount == 0){ 41 ltGini = 0.0f; 42 } 43 44 float gini = gteqGini * float(gteqCount) / samples; 45 gini += ltGini * float(ltCount) / samples; 46 47 return gini; 48 } 49