Python 读取 csv 文件并过滤数据

Question

I appologize if this was already answered before, but I checked I bunch of posts and just cannot understand what is wrong with my code. I'm trying to read a csv file in python (see bellow) and filter out rows of data by the value in the second column (angle). Then I want to create a new output file with filtered time and angle values. I only get the output file with headers written in.

csv file:

time,angle
0,56
1,89
2,112
3,189
4,122
5,123

Code:

import csv

#define the min and max value of angle
alpha_min = 110
alpha_max = 125

#read csv file and loop through with a filter
with open('test_csv.csv', 'r') as input_file:
    csv_reader = csv.reader(input_file)#, delimiter=',')
    #header = next(input_file).strip("\n").split(",")
    results = filter(lambda row: alpha_min<row[1]<alpha_max, csv_reader)

#create output file
with open('test_output_csv.csv', "w") as output_file:
    csv_writer = csv.writer(output_file, delimiter=',')
    csv_writer.writerow(header)
    for result in results:
        csv_writer.writerow(result)

Answer 1

I would suggest using pandas library for this workflow, which will be faster and more efficient than looping through each line of your csv file. Something like the below:

import pandas as pd

#define the min and max value of angle
alpha_min = 110
alpha_max = 125

# read input and filter angle data
df = pd.read_csv('test_csv.csv')
df = df[(df['angle'] < alpha_max) & (df['angle'] > alpha_min)]

# write output
df.to_csv('output.csv')

Answer 2

You can do

import csv

#define the min and max value of angle
alpha_min = 110
alpha_max = 125

#read csv file and loop through with a filter
with open('test_csv.csv', 'r') as input_file:
    csv_reader = csv.reader(input_file)#, delimiter=',')
    lines = [i for i in csv_reader]
    header = lines[0]
    results = filter(lambda row: alpha_min<int(row[1])<alpha_max, lines[1:])

#create output file
with open('test_output_csv.csv', "w", newline='') as output_file:
    csv_writer = csv.writer(output_file, delimiter=',')
    csv_writer.writerow(header) 
    csv_writer.writerows(results)

That will save to the file

time,angle
2,112
4,122
5,123

Answer 3

The fields of a csv row are strings so you need int(row[1]) to work correctly. I also recommend a list comprehension for the filtering, or pandas for speed. next(csv_reader) will read one row to capture the headers as well.

Note: use newline='' with the csv module as documented to avoid blank lines between each row .

import csv

alpha_min = 110
alpha_max = 125

with open('test.csv','r',newline='') as input_file:
    csv_reader = csv.reader(input_file)
    header = next(csv_reader)
    results = [row for row in csv_reader if alpha_min < int(row[1]) < alpha_max]

with open('output.csv','w',newline='') as output_file:
    csv_writer = csv.writer(output_file)
    csv_writer.writerow(header)
    csv_writer.writerows(results)