C++实现sklearn风格train_test_split函数遇返回值问题求解决

你的代码里有两个关键错误导致只有test矩阵有数据：

1. return train, test;是逗号表达式，C++中逗号表达式会返回最后一个值，所以函数实际只返回了test，train直接被丢弃。
2. train, test = train_test_split(df);同样是逗号表达式，只有test被赋值，train完全没拿到数据，所以输出里train.size()为0。

以下是几种可行的解决方案，按新手友好度排序：

方案1：使用引用参数作为输出（最常用） #

把train和test作为引用参数传给函数，函数直接在这两个容器里填充数据，不需要返回值。这种方式符合C++惯用写法，还能避免不必要的vector拷贝，效率更高。

修改后的代码：

#include <iostream> 
#include <cstdlib>
#include <time.h>
#include <vector>  

using namespace std;

// 将train和test设为引用参数，函数直接修改它们
void train_test_split(const vector<vector<float>>& df, vector<vector<float>>& train, vector<vector<float>>& test, float train_size = 0.8){
    train.clear();
    test.clear();
    for(size_t i = 0; i < df.size(); i++){
        // 生成0-99的随机数，支持更精细的比例（比如0.75、0.85）
        int x = rand() % 100; 
        if(x <= train_size * 100){
            train.push_back(df[i]);
        } 
        else{
            test.push_back(df[i]);
        }
    }
} 

int main(){
    vector<vector<float>> train;
    vector<vector<float>> test; 
    vector<vector<float>> df = {{1,2,3,4}, 
                               {5,6,7,8},
                               {9,10,11,12}};

    srand(time(NULL)); // srand只需调用一次，放在程序开头即可
    train_test_split(df, train, test); 

    cout << "training size: " << train.size() << ", test size: " << test.size() << endl; 
    return 0; 
}

注意点：

• 把df设为const &避免不必要的拷贝，提升效率。
• 不要把srand(time(NULL))放在循环或函数里，否则每次调用都会重置随机种子，导致随机数重复。

方案2：返回std::pair打包两个结果 #

如果偏好函数返回值的风格，可以用std::pair把两个vector打包返回，C++17及以上支持结构化绑定，能直接解构结果：

#include <iostream> 
#include <cstdlib>
#include <time.h>
#include <vector>  
#include <utility> // 引入std::pair

using namespace std;

pair<vector<vector<float>>, vector<vector<float>>> train_test_split(const vector<vector<float>>& df, float train_size = 0.8){
    vector<vector<float>> train;
    vector<vector<float>> test;
    for(size_t i = 0; i < df.size(); i++){
        int x = rand() % 100;
        if(x <= train_size * 100){
            train.push_back(df[i]);
        } else {
            test.push_back(df[i]);
        }
    }
    return {train, test}; // first是train，second是test
} 

int main(){
    vector<vector<float>> df = {{1,2,3,4}, 
                               {5,6,7,8},
                               {9,10,11,12}};

    srand(time(NULL));
    // C++17结构化绑定，直接解构pair到train和test
    auto [train, test] = train_test_split(df);

    cout << "training size: " << train.size() << ", test size: " << test.size() << endl; 
    return 0; 
}

如果编译器不支持C++17，可手动提取pair内容：

auto result = train_test_split(df);
vector<vector<float>> train = result.first;
vector<vector<float>> test = result.second;

方案3：自定义结构体（适合复杂返回场景） #

如果后续需要返回更多相关数据，自定义结构体可以让代码可读性更好：

#include <iostream> 
#include <cstdlib>
#include <time.h>
#include <vector>  

using namespace std;

// 自定义结构体存储拆分结果
struct TrainTestSplitResult {
    vector<vector<float>> train;
    vector<vector<float>> test;
};

TrainTestSplitResult train_test_split(const vector<vector<float>>& df, float train_size = 0.8){
    TrainTestSplitResult result;
    for(size_t i = 0; i < df.size(); i++){
        int x = rand() % 100;
        if(x <= train_size * 100){
            result.train.push_back(df[i]);
        } else {
            result.test.push_back(df[i]);
        }
    }
    return result;
} 

int main(){
    vector<vector<float>> df = {{1,2,3,4}, 
                               {5,6,7,8},
                               {9,10,11,12}};

    srand(time(NULL));
    auto result = train_test_split(df);

    cout << "training size: " << result.train.size() << ", test size: " << result.test.size() << endl; 
    return 0; 
}

总结 #

新手优先推荐方案1，它是C++处理多输出最常用的方式，效率高且容易理解；如果喜欢函数返回值的风格，方案2的pair足够简单；方案3适合需要返回更多关联数据的场景。

内容的提问来源于stack exchange，提问作者Olive Yew