如何使用切片获取指定行/列数据，实现按年份提取csv数据的Python函数

首先默认你的data.csv第一列为年份列，若年份列位置不同，修改代码中列索引数值即可（索引从0开始计数）。

核心原理说明 #

numpy加载csv得到的是二维数组，通用索引规则为数组[行选择规则, 列选择规则]：

• 单独冒号:表示选择该维度所有元素，比如data[:, 0]表示选择所有行的第0列（即全部年份值），data[2, :]表示选择第3行的所有列（即第3行全部数据）
• 可传入布尔数组做筛选，仅保留布尔值为True对应位置的元素

完整实现代码 #

import numpy as np

def get_data(year):
    # 修正原代码语法错误：文件名需要加引号，若文件不在运行目录需补全路径
    data = np.loadtxt("data.csv", delimiter=',', skiprows=1)
    year = int(year)
    # 提取所有行的第0列（年份列）
    all_years = data[:, 0]
    # 判断年份是否存在，不存在则抛出异常
    if year not in all_years:
        raise ValueError(f"未找到{year}年的对应数据")
    # 筛选年份匹配的行，返回该行全部数据
    # 若存在多条同一年份的数据，去掉末尾的[0]即可返回所有匹配行
    return data[all_years == year][0]

测试示例 #

假设你的data.csv内容如下：

年份,收入,支出
2021,1000,500
2022,1200,600
2023,1500,700

调用get_data(2022)会返回[2022, 1200, 600]，调用get_data(2020)会直接抛出ValueError异常。

内容的提问来源于stack exchange，提问作者Bayle