Share My C++ Solution

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;
    }
};
Comments (1)