Thanks very much for such topic and coding practice. Now I have better understanding of the decision trees.
struct Flower{
int plen;
string spec;
Flower() {}
Flower(int l, string s) {plen = l; spec = s;}
bool operator<(const Flower& other) const {return plen < other.plen;}
};
class Solution {
public:
double calculateMaxInfoGain(vector<double>& petal_length, vector<string>& species) {
int size = petal_length.size();
vector<Flower> flowers(size);
unordered_map<string, int> count;
for(int i=0; i<size; i++) {
flowers[i] = {petal_length[i], species[i]};
count[species[i]] += 1;
}
//calculate entropy of the whole vector
double H = getEntropy(count);
//calculate reduction of entropy
sort(flowers.begin(), flowers.end());
unordered_map<string, int> count1;
unordered_map<string, int>& count2 = count;
int size1 = 0, size2 = size;
double gain = DBL_MIN;
for(int i=0; i<size-1; i++) {
size1 += 1; size2 -= 1;
count1[flowers[i].spec] += 1;
count2[flowers[i].spec] -= 1;
if(count2[flowers[i].spec] == 0) count2.erase(flowers[i].spec);
double ent1 = getEntropy(count1);
double ent2 = getEntropy(count2);
double r = H - size1/(double)size*ent1 - size2/(double)size*ent2;
gain = max(gain, r);
}
return gain;
}
private:
double getEntropy(unordered_map<string, int>& emap) {
double H = 0.0; int total = 0;
for(auto& kv : emap) {total += kv.second;}
for(auto& kv : emap) {
double p = (double)kv.second / (double)total;
H += -p*log2(p);
}
return H;
}
};