| 25 | |
| 26 | |
| 27 | class TreeNode: |
| 28 | def __init__(self, depth=1, max_depth=None): |
| 29 | print('depth:', depth) |
| 30 | self.depth = depth |
| 31 | self.max_depth = max_depth |
| 32 | if self.max_depth is not None and self.max_depth < self.depth: |
| 33 | raise Exception("depth > max_depth") |
| 34 | |
| 35 | def fit(self, X, Y): |
| 36 | if len(Y) == 1 or len(set(Y)) == 1: |
| 37 | # base case, only 1 sample |
| 38 | # another base case |
| 39 | # this node only receives examples from 1 class |
| 40 | # we can't make a split |
| 41 | self.col = None |
| 42 | self.split = None |
| 43 | self.left = None |
| 44 | self.right = None |
| 45 | self.prediction = Y[0] |
| 46 | |
| 47 | else: |
| 48 | D = X.shape[1] |
| 49 | cols = range(D) |
| 50 | |
| 51 | max_ig = 0 |
| 52 | best_col = None |
| 53 | best_split = None |
| 54 | for col in cols: |
| 55 | ig, split = self.find_split(X, Y, col) |
| 56 | # print "ig:", ig |
| 57 | if ig > max_ig: |
| 58 | max_ig = ig |
| 59 | best_col = col |
| 60 | best_split = split |
| 61 | |
| 62 | if max_ig == 0: |
| 63 | # nothing we can do |
| 64 | # no further splits |
| 65 | self.col = None |
| 66 | self.split = None |
| 67 | self.left = None |
| 68 | self.right = None |
| 69 | self.prediction = np.round(Y.mean()) |
| 70 | else: |
| 71 | self.col = best_col |
| 72 | self.split = best_split |
| 73 | |
| 74 | if self.depth == self.max_depth: |
| 75 | self.left = None |
| 76 | self.right = None |
| 77 | self.prediction = [ |
| 78 | np.round(Y[X[:,best_col] < self.split].mean()), |
| 79 | np.round(Y[X[:,best_col] >= self.split].mean()), |
| 80 | ] |
| 81 | else: |
| 82 | # print "best split:", best_split |
| 83 | left_idx = (X[:,best_col] < best_split) |
| 84 | # print "left_idx.shape:", left_idx.shape, "len(X):", len(X) |