python - 如何将数据随机拆分为训练集和测试集？

Question

我有一个大数据集，想把它分成训练集(50%)和测试集(50%)。

假设我有 100 个示例存储输入文件，每一行包含一个示例。我需要选择 50 行作为训练集和 50 行测试集。

我的想法是首先生成一个长度为 100 的随机列表(值范围从 1 到 100)，然后使用前 50 个元素作为 50 个训练示例的行号。测试集也是如此。

这在 Matlab 中很容易实现

fid=fopen(datafile);
C = textscan(fid, '%s','delimiter', '\n');
plist=randperm(100);
for i=1:50
    trainstring = C{plist(i)};
    fprintf(train_file,trainstring);
end
for i=51:100
    teststring = C{plist(i)};
    fprintf(test_file,teststring);
end

但是我怎样才能在 Python 中完成这个功能呢？我是 Python 新手，不知道是否可以将整个文件读入一个数组，然后选择某些行。

Answer 1

这可以在 Python 中使用列表类似地完成，(请注意，整个列表在原地打乱)。

import random

with open("datafile.txt", "rb") as f:
    data = f.read().split('\n')

random.shuffle(data)

train_data = data[:50]
test_data = data[50:]