C++代码实现MATLAB中的fitrtree函数功能

// fitrtree.cpp
// 编译: g++ -std=c++17 -O2 fitrtree.cpp -o fitrtree

#include <vector>
#include <memory>
#include <string>
#include <limits>
#include <algorithm>
#include <numeric>
#include <cmath>
#include <stdexcept>
#include <iostream>
#include <iomanip>
#include <map>

namespace mltree {

// ============================================================
// RegressionTree 节点 ------ 对齐 MATLAB RegressionTree 属性
//   CutPoint          : 连续型分裂的切点
//   CutPredictorIndex : 分裂使用的预测变量索引 (0-based)
//   NodeSize          : 节点样本数
//   NodeMean          : 节点响应均值 (叶节点预测值)
//   NodeMSE           : 节点 MSE
//   IsBranchNode      : 是否为分支节点
// ============================================================
struct Node {
    int    cut_predictor_index = -1;   // -1 表示叶节点
    double cut_point           = 0.0;  // 切点 u，规则: x_j < u -> 左
    double node_mean           = 0.0;  // 叶节点预测值
    double node_mse            = 0.0;
    int    node_size           = 0;
    std::unique_ptr<Node> left;
    std::unique_ptr<Node> right;

    bool is_branch() const { return cut_predictor_index >= 0; }
    bool is_leaf()   const { return cut_predictor_index < 0; }
};

// ============================================================
// fitrtree 参数 ------ 对齐 MATLAB Name-Value 对
//   MinParentSize : 默认 10
//   MinLeafSize   : 默认 1
//   MaxNumSplits  : 默认 size(X,1)-1
//   MergeLeaves   : 默认 'on'
//   Prune         : 默认 'on'
// ============================================================
struct FitRTreeParams {
    int    MinParentSize = 10;
    int    MinLeafSize   = 1;
    int    MaxNumSplits  = -1;   // -1 表示用 n-1
    bool   MergeLeaves   = true;
    bool   Prune         = true;
    double PruneAlpha    = 0.0;
};

// ============================================================
// 回归树 ------ 对应 MATLAB RegressionTree 类
// ============================================================
class RegressionTree {
public:
    RegressionTree() = default;

    // 对应 tree = fitrtree(X, Y, Name, Value)
    void fit(const std::vector<std::vector<double>>& X,
             const std::vector<double>& Y,
             const FitRTreeParams& params = FitRTreeParams()) {

        if (X.empty() || Y.empty() || X.size() != Y.size())
            throw std::invalid_argument("X and Y must have the same non-zero size");
        if (X[0].empty())
            throw std::invalid_argument("X must have at least one feature");

        params_ = params;
        n_features_ = static_cast<int>(X[0].size());
        n_samples_  = static_cast<int>(X.size());

        // MATLAB 默认 MaxNumSplits = n - 1
        if (params_.MaxNumSplits < 0)
            params_.MaxNumSplits = n_samples_ - 1;

        std::vector<int> indices(n_samples_);
        std::iota(indices.begin(), indices.end(), 0);

        // 步骤 1: 生长完整树
        root_ = build_tree(X, Y, indices, 0);

        // 步骤 2: MergeLeaves='on' ------ 合并同父叶节点 (MATLAB 默认)
        if (params_.MergeLeaves)
            merge_leaves(root_.get());

        // 步骤 3: Prune='on' ------ 代价复杂度剪枝 (MATLAB 默认)
        if (params_.Prune)
            prune_cost_complexity(root_.get(), params_.PruneAlpha);
    }

    // 对应 Yfit = predict(Mdl, X)
    // MATLAB 规则: 若 x_j < u 走左枝，否则走右枝
    double predict(const std::vector<double>& x) const {
        if (!root_) throw std::runtime_error("Model not fitted");
        const Node* node = root_.get();
        while (node->is_branch()) {
            node = (x[node->cut_predictor_index] < node->cut_point)
                   ? node->left.get() : node->right.get();
        }
        return node->node_mean;
    }

    std::vector<double> predict(const std::vector<std::vector<double>>& X) const {
        std::vector<double> result;
        result.reserve(X.size());
        for (const auto& x : X) result.push_back(predict(x));
        return result;
    }

    int num_nodes() const {
        int count = 0;
        count_nodes(root_.get(), count);
        return count;
    }

    int num_leaves() const { return count_leaves(root_.get()); }

    // 对应 resubLoss(Mdl)
    double resub_loss(const std::vector<std::vector<double>>& X,
                      const std::vector<double>& Y) const {
        double sse = 0.0;
        for (size_t i = 0; i < X.size(); ++i) {
            double d = Y[i] - predict(X[i]);
            sse += d * d;
        }
        return sse / Y.size();
    }

private:
    FitRTreeParams params_;
    std::unique_ptr<Node> root_;
    int n_features_ = 0;
    int n_samples_  = 0;

    // ------------------------------------------------------------
    // 递归生长树 ------ 标准 CART
    // ------------------------------------------------------------
    std::unique_ptr<Node> build_tree(
        const std::vector<std::vector<double>>& X,
        const std::vector<double>& Y,
        const std::vector<int>& indices,
        int splits_so_far) {

        auto node = std::make_unique<Node>();
        node->node_size = static_cast<int>(indices.size());

        double sum = 0.0;
        for (int i : indices) sum += Y[i];
        node->node_mean = sum / indices.size();

        double sse = 0.0;
        for (int i : indices) {
            double d = Y[i] - node->node_mean;
            sse += d * d;
        }
        node->node_mse = sse / indices.size();

        // 停止条件 (对齐 MATLAB Growing Decision Trees)
        if (node->node_size < params_.MinParentSize) return node;
        if (splits_so_far >= params_.MaxNumSplits) return node;
        if (sse < 1e-12) return node;

        SplitResult best = find_best_split(X, Y, indices, sse);
        if (best.feature_index < 0) return node;

        node->cut_predictor_index = best.feature_index;
        node->cut_point           = best.threshold;

        std::vector<int> left_idx, right_idx;
        left_idx.reserve(best.n_left);
        right_idx.reserve(node->node_size - best.n_left);
        // 使用 < 而非 <= ------ 与 MATLAB 的 x_j < u 规则一致
        for (int i : indices) {
            if (X[i][best.feature_index] < best.threshold)
                left_idx.push_back(i);
            else
                right_idx.push_back(i);
        }

        node->left  = build_tree(X, Y, left_idx,  splits_so_far + 1);
        node->right = build_tree(X, Y, right_idx, splits_so_far + 1);

        return node;
    }

    struct SplitResult {
        int    feature_index = -1;
        double threshold = 0.0;
        double mse_improvement = 0.0;
        int    n_left = 0;
    };

    // ------------------------------------------------------------
    // 最优分裂搜索
    // ------------------------------------------------------------
    SplitResult find_best_split(
        const std::vector<std::vector<double>>& X,
        const std::vector<double>& Y,
        const std::vector<int>& indices,
        double parent_sse) {

        SplitResult best;

        for (int f = 0; f < n_features_; ++f) {
            std::vector<int> sorted_idx = indices;
            std::sort(sorted_idx.begin(), sorted_idx.end(),
                      [&](int a, int b) { return X[a][f] < X[b][f]; });

            double total_sum = 0.0, total_sq_sum = 0.0;
            for (int i : indices) {
                total_sum    += Y[i];
                total_sq_sum += Y[i] * Y[i];
            }

            double left_sum = 0.0, left_sq_sum = 0.0;
            int    left_n   = 0;

            for (int k = 0; k < static_cast<int>(sorted_idx.size()) - 1; ++k) {
                int i = sorted_idx[k];
                left_sum    += Y[i];
                left_sq_sum += Y[i] * Y[i];
                left_n++;

                if (X[sorted_idx[k]][f] == X[sorted_idx[k + 1]][f]) continue;

                int right_n = static_cast<int>(indices.size()) - left_n;
                if (left_n < params_.MinLeafSize || right_n < params_.MinLeafSize)
                    continue;

                double left_sse  = left_sq_sum - left_sum * left_sum / left_n;
                double right_sse = (total_sq_sum - left_sq_sum)
                                   - (total_sum - left_sum) * (total_sum - left_sum) / right_n;
                if (left_sse < 0)  left_sse  = 0;
                if (right_sse < 0) right_sse = 0;

                double improvement = parent_sse - (left_sse + right_sse);

                if (improvement > best.mse_improvement) {
                    best.feature_index   = f;
                    best.threshold       = (X[sorted_idx[k]][f]
                                          + X[sorted_idx[k + 1]][f]) / 2.0;
                    best.mse_improvement = improvement;
                    best.n_left          = left_n;
                }
            }
        }
        return best;
    }

    // ------------------------------------------------------------
    // MergeLeaves ------ 对应 MATLAB 默认 MergeLeaves='on'
    // ------------------------------------------------------------
    void merge_leaves(Node* node) {
        if (!node || node->is_leaf()) return;

        merge_leaves(node->left.get());
        merge_leaves(node->right.get());

        if (node->left->is_leaf() && node->right->is_leaf()) {
            double left_mse  = node->left->node_mse  * node->left->node_size;
            double right_mse = node->right->node_mse * node->right->node_size;
            double parent_mse = node->node_mse * node->node_size;

            if (left_mse + right_mse >= parent_mse) {
                node->cut_predictor_index = -1;
                node->cut_point = 0.0;
                node->left.reset();
                node->right.reset();
            }
        }
    }

    // ------------------------------------------------------------
    // 代价复杂度剪枝 ------ 对应 Prune='on'
    // ------------------------------------------------------------
    void prune_cost_complexity(Node* node, double alpha) {
        if (!node || node->is_leaf()) return;

        prune_cost_complexity(node->left.get(),  alpha);
        prune_cost_complexity(node->right.get(), alpha);

        if (node->left->is_leaf() && node->right->is_leaf()) {
            double sse_before = node->left->node_mse  * node->left->node_size
                              + node->right->node_mse * node->right->node_size
                              + alpha * 2.0;
            double sse_after  = node->node_mse * node->node_size
                              + alpha * 1.0;

            if (sse_after <= sse_before) {
                node->cut_predictor_index = -1;
                node->cut_point = 0.0;
                node->left.reset();
                node->right.reset();
            }
        }
    }

    int count_leaves(const Node* node) const {
        if (!node) return 0;
        if (node->is_leaf()) return 1;
        return count_leaves(node->left.get()) + count_leaves(node->right.get());
    }

    void count_nodes(const Node* node, int& count) const {
        if (!node) return;
        count++;
        count_nodes(node->left.get(),  count);
        count_nodes(node->right.get(), count);
    }
};

} // namespace mltree

// ============================================================
// 使用示例
// ============================================================
int main() {
    std::vector<std::vector<double>> X = {
        {1.0}, {2.0}, {3.0}, {4.0}, {5.0}, {6.0}, {7.0}, {8.0}
    };
    std::vector<double> Y = {1.0, 1.5, 2.0, 3.5, 5.0, 5.5, 7.0, 8.5};

    // --------------------------------------------------------
    // 场景 A: 默认参数 (MinParentSize=10)
    //   8 样本 < 10 -> 根节点即为叶节点
    // --------------------------------------------------------
    {
        mltree::FitRTreeParams p_default;
        mltree::RegressionTree tree;
        tree.fit(X, Y, p_default);

        std::cout << "=== Scenario A: MATLAB defaults ===\n";
        std::cout << "NumNodes  : " << tree.num_nodes()  << "\n";
        std::cout << "NumLeaves : " << tree.num_leaves() << "\n";
        std::cout << "RootMean  : " << tree.predict(std::vector<double>{0.0}) << "\n";
        std::cout << "ResubLoss : " << tree.resub_loss(X, Y) << "\n\n";
    }

    // --------------------------------------------------------
    // 场景 B: MinParentSize=4 (MATLAB 文档示例)
    // --------------------------------------------------------
    {
        mltree::FitRTreeParams p;
        p.MinParentSize = 4;
        p.MinLeafSize   = 1;
        p.MaxNumSplits  = -1;

        mltree::RegressionTree tree;
        tree.fit(X, Y, p);

        std::cout << "=== Scenario B: MinParentSize=4 ===\n";
        std::cout << "NumNodes  : " << tree.num_nodes()  << "\n";
        std::cout << "NumLeaves : " << tree.num_leaves() << "\n";

        std::vector<double> Xtest = {1.5, 4.5, 7.5};
        for (double v : Xtest) {
            std::cout << "predict(" << v << ") = "
                      << std::setprecision(4) << tree.predict(std::vector<double>{v}) << "\n";
        }
        std::cout << "ResubLoss : " << tree.resub_loss(X, Y) << "\n";
    }

    return 0;
}