| 119 | return max_ig, best_split |
| 120 | |
| 121 | def information_gain(self, x, y, split): |
| 122 | # assume classes are 0 and 1 |
| 123 | # print "split:", split |
| 124 | y0 = y[x < split] |
| 125 | y1 = y[x >= split] |
| 126 | N = len(y) |
| 127 | y0len = len(y0) |
| 128 | if y0len == 0 or y0len == N: |
| 129 | return 0 |
| 130 | p0 = float(len(y0)) / N |
| 131 | p1 = 1 - p0 #float(len(y1)) / N |
| 132 | # print "entropy(y):", entropy(y) |
| 133 | # print "p0:", p0 |
| 134 | # print "entropy(y0):", entropy(y0) |
| 135 | # print "p1:", p1 |
| 136 | # print "entropy(y1):", entropy(y1) |
| 137 | return entropy(y) - p0*entropy(y0) - p1*entropy(y1) |
| 138 | |
| 139 | def predict_one(self, x): |
| 140 | # use "is not None" because 0 means False |