python - 使用 CSV 文件创建维恩图的 Python 索引错误

Question

我正在尝试使用我创建的 CSV 文件中的信息创建维恩图。但是，我在代码的第 13 行收到一个索引错误，它说列表索引超出范围。我希望看看这里是否有人对为什么会这样有任何想法。这是我用来尝试创建维恩图的代码：

from matplotlib_venn import venn2
import matplotlib.pyplot as plt
import csv
from sympy import FiniteSet

def get_Sports(file_name):
    football=[]
    others=[]
    with open(file_name) as file:
        reader = csv.reader(file)
        next(reader)
        for row in reader:
            if row[1] == 1:
                football.append(row[0])
            if row[2] == 1:
                others.append(row[0])
    return football, others

def plot_venn(f, o):
    venn2(subsets=(f, o))
    plt.show()

if __name__ == '__main__':
    file = input('Input the file path with the list of student ids, whether they play football, and whether they play other sports: ')
    football, others = get_Sports(file)
    football = FiniteSet(*football)
    others = FiniteSet(*others)
    plot_venn(football, others)

这是我正在使用的 CSV 文件：

ID, Football, Other Sports
1, 1, 0
2, 0, 1
3, 0, 0
4, 1, 1
5, 1, 0
6, 0, 0
7, 0, 0
8, 1, 1
9, 1, 1
10, 1, 0

标签： pythoncsv

Answer 1

您的逻辑中有一个错误，因为当您读取 csv 时，它将值存储为字符串。所以row[1] == 1将是False因为row[1]是'1'。

为了使您的代码正常工作，您的功能需要是：

def get_Sports(file_name):
    football=[]
    others=[]
    with open(file_name) as file:
        reader = csv.reader(file)
        next(reader)
        for row in reader:
            if row[1].strip() == '1':
                football.append(row[0])
            if row[2].strip() == '1':
                others.append(row[0])
    return football, others

我个人pandas比较喜欢csv. 所以我和你做了同样的事情，只是使用 pandas，然后使用逻辑最终得到你想要的。如果您的 csv 文件很大，这将是一种更好/更快的方法：

from matplotlib_venn import venn2
import matplotlib.pyplot as plt
import pandas as pd
from sympy import FiniteSet

def get_Sports(file_name):
    df = pd.read_csv(file_name)
    football = list(df[df.iloc[:,1] == 1].iloc[:,0])
    others = list(df[df.iloc[:,-1] == 1].iloc[:,0])
    
    return football, others

def plot_venn(f, o):
    venn2(subsets=(f, o), set_labels = ('Football', 'Other Sports'))
    plt.show()

if __name__ == '__main__':
    file = input('Input the file path with the list of student ids, whether they play football, and whether they play other sports: ')
    football, others = get_Sports(file)
    football = FiniteSet(*football)
    others = FiniteSet(*others)
    plot_venn(football, others)

输出：