pandas如何根据用户输入的指定列计算该列的总和

原代码存在的问题 #

• df.at[user_input_grade]用法错误：at方法用于定位单个单元格，需要同时传入行索引和列索引，无法直接选中整列，选择指定列应该用df[列名]的语法
• 没有做输入合法性校验：如果用户输入的列名不存在于数据表中，代码会直接抛出异常
• 存在冗余的pandas导入语句：开头导入一次即可，函数内部不需要重复导入
• input(str("enter the grade: "))写法冗余：input方法本身返回字符串，不需要额外用str()转换

修正后的可运行代码 #

import pandas as pd

def get_additonal_data():
    file_path = '/Users/CIS5357/Assignments/'
    data = file_path + 'grades.csv'
    df_grades = pd.read_csv(data)
    print("读取到的成绩数据如下：")
    print(df_grades)
    return df_grades

def do_functions(df_grades):
    # 先输出所有可用列名，提示用户输入
    print("\n数据表中包含的列有：", list(df_grades.columns))
    user_input_col = input("请输入要计算总和的列名：")
    # 先校验列名是否存在
    if user_input_col not in df_grades.columns:
        print(f"错误：不存在名为【{user_input_col}】的列")
        return
    # 计算指定列的总和
    col_sum = df_grades[user_input_col].sum()
    print(f"列【{user_input_col}】的总和为：{col_sum}")

if __name__ == "__main__":
    df_grades = get_additonal_data()
    do_functions(df_grades)

核心实现逻辑说明 #

• 接收用户输入的列名后，先通过df.columns判断列名是否合法，避免运行报错
• 选中指定列用df[列名]的语法，直接调用.sum()方法即可得到该列的总和，pandas会自动跳过空值计算
• 提前输出数据表所有可用列名，降低用户输入错误的概率

内容的提问来源于stack exchange，提问作者Namitha Subbarama Rao