基于奇偶年份的二维数组训练测试集拆分函数实现求助

需求回顾 #

输入为二维numpy数组，每行格式为[年份, 对应数值]，需返回格式为((X_train, y_train), (X_test, y_test))的结果：

• (X_train, y_train)：包含所有偶数年份及其对应数值
• (X_test, y_test)：包含所有奇数年份及其对应数值

原代码问题分析 #

1. 错误使用train_test_split：需求不需要随机拆分数据集，而是直接将全量偶数年份作为训练集、全量奇数年份作为测试集
2. 奇偶行索引搞反：输入数组第0行是奇数年份（1961）、第1行是偶数年份（1962），原代码用arr[::2]取奇数年份行却误当作偶数处理
3. 数据分配错误：转置后错误将年份和数值的位置混淆，导致输出内容完全错位

修正后的代码 #

import numpy as np

def feat_resp_split(arr):
    # 筛选偶数年份的行
    even_rows = arr[arr[:, 0] % 2 == 0]
    # 筛选奇数年份的行
    odd_rows = arr[arr[:, 0] % 2 == 1]
    
    # 拆分年份（X）和对应数值（y）
    X_train, y_train = even_rows[:, 0], even_rows[:, 1]
    X_test, y_test = odd_rows[:, 0], odd_rows[:, 1]
    
    return (X_train, y_train), (X_test, y_test)

关键说明 #

• 通过年份数值直接判断奇偶，不依赖输入数组的行顺序，鲁棒性更强
• 明确拆分每行第一列（年份）为X、第二列（数值）为y，完全匹配预期输出格式
• 移除不必要的train_test_split操作，直接返回全量奇偶年份数据集

内容的提问来源于stack exchange，提问作者Jean Swiegers