请求协助修正Python人口增长率计算函数（返回Numpy数组）

需求说明 #

• 输入：population_df（Pandas DataFrame）和country_code字符串
• 计算逻辑：增长率公式为 Growth_rate=(current_year_population−previous_year_population)/previous_year_population
• 限制：仅计算1961年及以后的增长率
• 输出：形状为(?,2)的二维Numpy数组，每一行对应[年份, 增长率]

原代码 #

def pop_growth_by_country_year(df,country_code):
    country_data = df.loc[country_code]
    for columnName, columnData in country_data.iteritems():
        country_data = ((country_data[columnData] - country_data[columnData-1]) / country_data[columnData-1])
    output = country_data.reset_index().to_numpy().reshape(-1, 2)
    return output

修正后的代码 #

import pandas as pd
import numpy as np

def pop_growth_by_country_year(df, country_code):
    # 获取指定国家的人口数据，按年份排序
    country_data = df.loc[country_code].sort_index()
    
    # 利用shift方法计算相邻年份的增长率
    growth_rates = (country_data - country_data.shift(1)) / country_data.shift(1)
    
    # 过滤1961年及以后的数据，剔除无上年数据的无效值
    growth_rates = growth_rates.loc[growth_rates.index >= 1961].dropna()
    
    # 组合年份与增长率，转换为要求的二维Numpy数组
    output = np.column_stack((growth_rates.index.values.astype(float), growth_rates.values.astype(float)))
    
    return output

关键修正点 #

• 修复遍历逻辑：原循环逻辑混乱，改用Pandas的shift()方法高效计算相邻年份差值，避免手动循环错误
• 新增年份过滤：通过loc[growth_rates.index >= 1961]严格限制仅保留1961年及以后的增长率数据
• 优化数组生成：用column_stack直接组合年份索引与增长率值，生成符合(?,2)形状的数组，无需多余的reshape操作
• 处理无效值：用dropna()移除无上年数据的第一年（1960年及更早的无效记录）

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