将一个CSV格式的文件分割成两个CSV文件

序言

本项目可以按照比例将一个csv文件分割成两个csv文件，效果是：在C:\algo_file文件夹下，将该文件夹下的data.csv文件分成train.csv和vali.csv

完整代码

定义split_csv函数：

import csv
import os


def split_csv(path, total_len, per):

    # 如果train.csv和vali.csv存在就删除
    if os.path.exists('C:\\algo_file\\train.csv'):
        os.remove('C:\\algo_file\\train.csv')
    if os.path.exists('C:\\algo_file\\vali.csv'):
        os.remove('C:\\algo_file\\vali.csv')

    with open(path, 'r', newline='') as file:
        csvreader = csv.reader(file)
        i = 0
        for row in csvreader:

            if i < round(total_len * per/100):
                # train.csv存放路径
                csv_path = os.path.join("C:\\algo_file", 'train.csv')
                print(csv_path)
                # 不存在此文件的时候，就创建
                if not os.path.exists(csv_path):
                    with open(csv_path, 'w', newline='') as file:
                        csvwriter = csv.writer(file)
                        csvwriter.writerow(row)
                    i += 1
                # 存在的时候就往里面添加
                else:
                    with open(csv_path, 'a', newline='') as file:
                        csvwriter = csv.writer(file)
                        csvwriter.writerow(row)
                    i += 1
            elif (i >= round(total_len * per/100)) and (i < total_len):
            	# vali.csv存放路径
                csv_path = os.path.join("C:\\algo_file", 'vali.csv')
                print(csv_path)
                # 不存在此文件的时候，就创建
                if not os.path.exists(csv_path):
                    with open(csv_path, 'w', newline='') as file:
                        csvwriter = csv.writer(file)
                        csvwriter.writerow(row)
                    i += 1
                # 存在的时候就往里面添加
                else:
                    with open(csv_path, 'a', newline='') as file:
                        csvwriter = csv.writer(file)
                        csvwriter.writerow(row)
                    i += 1
            else:
                break

    print("训练集和验证集分离成功")
    return

调用上述函数：

if __name__ == '__main__':

    path = 'C:\\algo_file\\data.csv'
    total_len = len(open(path, 'r').readlines())# csv文件行数
    per = 80 # 分割比例%

    split_csv(path, total_len, per)

按照训练集占80%（验证集20%）比例，对C:\algo_file文件夹下的data.csv进行分割，在该文件下得到train.csv 和 vali.csv。

最后

本项目只是以C:\algo_file文件夹为例，实际上data.csv所在路径，train.csv所在路径，vali.csv所在路径包括文件名都可以更改。