自学教程：回归树的原理及Python实现

class Node:
    def __init__(self, score=None):
        self.score = score
        self.left = None
        self.right = None
        self.feature = None
        self.split = None

class RegressionTree(object):
    def __init__(self):
        self.root = Node()
        self.depth = 1
        self._rules = None

@staticmethod
def _get_split_mse(col: array, score: array, split: float):
    # Split score.
    score_left = score[col < split]
    score_right = score[col >= split]

    # Calculate the means of score.
    avg_left = score_left.mean()
    avg_right = score_right.mean()

    # Calculate the mse of score.
    mse = (((score_left - avg_left) ** 2).sum() +
            ((score_right - avg_right) ** 2).sum()) / len(score)

    return mse, avg_left, avg_right

def _choose_split(self, col: array, score: array):
    # Feature cannot be splitted if there's only one unique element.
    unique = set(col)
    if len(unique) == 1:
        return None, None, None, None
    # In case of empty split
    unique.remove(min(unique))

    # Get split point which has min mse
    ite = map(lambda x: (*self._get_split_mse(col, score, x), x), unique)
    mse, avg_left, avg_right, split = min(ite, key=lambda x: x[0])

    return mse, avg_left, avg_right, split

def _choose_feature(self, data: array, score: array):
    # Compare the mse of each feature and choose best one.
    ite = map(lambda x: (*self._choose_split(
        data[:, x], score), x), range(data.shape[1]))
    ite = filter(lambda x: x[0] is not None, ite)

    # Terminate if no feature can be splitted
    return min(ite, default=None, key=lambda x: x[0])

@staticmethod
def _expr2literal(expr):
    feature, operation, split = expr
    operation = ">=" if operation == 1 else "<"
    return "Feature%d %s %.4f" % (feature, operation, split)

def _get_rules(self):
    que = [[self.root, []]]
    self._rules = []
    # Breadth-First Search
    while que:
        node, exprs = que.pop(0)
        # Generate a rule when the current node is leaf node
        if not(node.left or node.right):
            # Convert expression to text
            literals = list(map(self._expr2literal, exprs))
            self._rules.append([literals, node.score])
        # Expand when the current node has left child
        if node.left:
            rule_left = copy(exprs)
            rule_left.append([node.feature, -1, node.split])
            que.append([node.left, rule_left])
        # Expand when the current node has right child
        if node.right:
            rule_right = copy(exprs)
            rule_right.append([node.feature, 1, node.split])
            que.append([node.right, rule_right])

def fit(self, data: array, score: array, max_depth=5, min_samples_split=2):
    # Initialize with depth, node, indexes
    self.root.score = score.mean()
    que = [(self.depth + 1, self.root, data, score)]
    # Breadth-First Search
    while que:
        depth, node, _data, _score = que.pop(0)
        # Terminate loop if tree depth is more than max_depth
        if depth > max_depth:
            depth -= 1
            break
        # Stop split when number of node samples is less than
        # min_samples_split or Node is 100% pure.
        if len(_score) < min_samples_split or all(_score == score[0]):
            continue
        # Stop split if no feature has more than 2 unique elements
        split_ret = self._choose_feature(_data, _score)
        if split_ret is None:
            continue
        # Split
        _, avg_left, avg_right, split, feature = split_ret
        # Update properties of current node
        node.feature = feature
        node.split = split
        node.left = Node(avg_left)
        node.right = Node(avg_right)
        # Put children of current node in que
        idx_left = (_data[:, feature] < split)
        idx_right = (_data[:, feature] >= split)
        que.append(
            (depth + 1, node.left, _data[idx_left], _score[idx_left]))
        que.append(
            (depth + 1, node.right, _data[idx_right], _score[idx_right]))
    # Update tree depth and rules
    self.depth = depth
    self._get_rules()

def __str__(self):
    ret = []
    for i, rule in enumerate(self._rules):
        literals, score = rule
        ret.append("Rule %d: " % i + ' | '.join(literals) + ' => y_hat %.4f' % score)
    return "\n".join(ret)

def _predict(self, row: array)->float:
    node = self.root
    while node.left and node.right:
        if row[node.feature] < node.split:
            node = node.left
        else:
            node = node.right
    return node.score

def predict(self, data: array)->array:
    return np.apply_along_axis(self._predict, 1, data)

def main():
    print("Tesing the performance of RegressionTree...")
    # Load data
    data, score = load_boston_house_prices()
    # Split data randomly, train set rate 70%
    data_train, data_test, score_train, score_test = train_test_split(
        data, score, random_state=200)
    # Train model
    reg = RegressionTree()
    reg.fit(data=data_train, score=score_train, max_depth=5)
    # Show rules
    print(reg)
    # Model evaluation
    get_r2(reg, data_test, score_test)