| 94 | self.right.fit(Xright, Yright) |
| 95 | |
| 96 | def find_split(self, X, Y, col): |
| 97 | # print "finding split for col:", col |
| 98 | x_values = X[:, col] |
| 99 | sort_idx = np.argsort(x_values) |
| 100 | x_values = x_values[sort_idx] |
| 101 | y_values = Y[sort_idx] |
| 102 | |
| 103 | # Note: optimal split is the midpoint between 2 points |
| 104 | # Note: optimal split is only on the boundaries between 2 classes |
| 105 | |
| 106 | # if boundaries[i] is true |
| 107 | # then y_values[i] != y_values[i+1] |
| 108 | # nonzero() gives us indices where arg is true |
| 109 | # but for some reason it returns a tuple of size 1 |
| 110 | boundaries = np.nonzero(y_values[:-1] != y_values[1:])[0] |
| 111 | best_split = None |
| 112 | max_ig = 0 |
| 113 | for b in boundaries: |
| 114 | split = (x_values[b] + x_values[b+1]) / 2 |
| 115 | ig = self.information_gain(x_values, y_values, split) |
| 116 | if ig > max_ig: |
| 117 | max_ig = ig |
| 118 | best_split = split |
| 119 | return max_ig, best_split |
| 120 | |
| 121 | def information_gain(self, x, y, split): |
| 122 | # assume classes are 0 and 1 |