数据分析中一段代码的优化过程

这段代码的目标是将新冠疫情时期，同一个国家的新冠累积确诊数据和累积康复数据做减法运算，以得到当前这一国家的当前的确诊数据，这样更能反映国家的医疗水平。

其中，confirmed_df是确诊数据，recovered_df是康复数据。两个DataFrame对象均有相同的columns如下:

Province/State,Country/Region,Lat,Long,1/22/20-3/9/23

我的基本想法就是先以国家进行分组，然后以confirmed_df的Country/Region为基准, 让recovered_df与其逐一对应。随后将两个DataFrame对象转化成numpy的ndarray多维数组对象，做矩阵的减法运算。

于是，第一版代码就出现了:

class DataAnalysis(object):
    def __init__(self):
        ...

    def get_current_confirmer_data(self):
        group_by_param = {"by": "Country/Region", "as_index": False}
        datetime_col = self.confirmed_df.columns[4:]
        sum_confirmed = self.confirmed_df.groupby(**group_by_param)[datetime_col].sum()
        sum_recovered = self.recovered_df.groupby(**group_by_param)[datetime_col].sum()
        
        confirmed_matrix = sum_confirmed[datetime_col].to_numpy()
        recovered_matrix = np.zeros_like(confirmed_matrix)

        for index, country in enumerate(sum_confirmed["Country/Region"]):
            match_row: pd.DataFrame = sum_recovered[sum_recovered["Country/Region"]==country]
            recovered_array: np.ndarray = match_row.iloc[:,1:].to_numpy()
            if recovered_array.size != 0:
                recovered_matrix[index] = recovered_array
            
        result = confirmed_matrix - recovered_matrix
        current_confirmed_df = pd.DataFrame(result, columns=datetime_col)
        current_confirmed_df["Country/Region"] = sum_confirmed["Country/Region"]
        return current_confirmed_df

可以看到，在代码中，使用了一个显式的for循环。在单独拿出sum_confirmed的Country/Region列后，我逐个遍历这些国家名，并在sum_recovered中寻找对应，从而让sum_recovered的每一行数据都能和sum_confirmed对应上。

这一思路简单质朴，但是使用Python的for循环来做数据分析往往都不见得高效。我猜想numpy或pandas肯定有方法能做类似需求。

我随后问ai让其优化代码，然而其优化思路不是很好:
 

def get_current_confirmed_data(self):
    # 用确诊数据与康复数据做减法
    group_by_param = {"by": "Country/Region", "as_index": False}
    datetime_col = self.confirmed_df.columns[4:]
    sum_confirmed = self.confirmed_df.groupby(**group_by_param)[datetime_col].sum()
    sum_recovered = self.recovered_df.groupby(**group_by_param)[datetime_col].sum()
    confirmed_matrix = sum_confirmed[datetime_col].to_numpy()
    recovered_matrix = np.zeros_like(confirmed_matrix)

    # 构建一个国家到索引的映射
    country_to_index = {country: idx for idx, country in enumerate(countries)}

    # 遍历康复数据，填充 recovered_matrix
    for _, row in sum_recovered.iterrows():
        country_idx = country_to_index.get(row["Country/Region"])
        if country_idx is not None:
            recovered_matrix[country_idx, :] = row[date_col_formatted].to_numpy()
    
    
    result = confirmed_matrix - recovered_matrix
    current_confirmed_df = pd.DataFrame(result, columns=datetime_col)
    current_confirmed_df["Country/Region"] = sum_confirmed["Country/Region"]
    return current_confirmed_df

这一思路似乎仅仅是为了改良原先的这句代码:

match_row: pd.DataFrame = sum_recovered[sum_recovered["Country/Region"]==country]

因为这句代码使用了布尔索引的匹配机制，逐一遍历查找对应的Country/Region并不高效。使用字典可以建立一种索引机制，加快匹配的速度。然后，由于本身数据量不大，因为全世界的国家数量本身就不多，遍历查找的效率不见得比索引慢。事实上，我在使用time模块测试了代码性能后发现，这段代码的性能甚至比原来更差，慢了大约0.03s左右。

于是我就去翻阅pandas的中文文档了。我直觉认为这就是和索引相关的内容，文档网址:

Python Pandas 极客教程

文档中, 重新索引的名字吸引了我，当我点开看其描述，发现高度符合我的需求。

于是，代码进一步优化，改写如下:

def tmp(self):
    group_by_param = {"by": "Country/Region", "as_index": True}
    datetime_col = self.confirmed_df.columns[4:]

    sum_confirmed = self.confirmed_df.groupby(**group_by_param)[datetime_col].sum()
    sum_recovered = self.recovered_df.groupby(**group_by_param)[datetime_col].sum()

    confirmed_matrix = sum_confirmed.to_numpy()
    recovered_matrix = sum_recovered.reindex(sum_confirmed.index).to_numpy()
        
    np.nan_to_num(recovered_matrix, copy=False)

    result = confirmed_matrix - recovered_matrix
    current_confirmed_df = pd.DataFrame(result, columns=datetime_col)
    current_confirmed_df["Country/Region"] = sum_confirmed.index

    return current_confirmed_df

这次的代码，显式的避免了for循环，使用pandas内置的reindex方法进行重新索引，同时使用numpy的nan_to_num对没有匹配上数据行形成的NaN值进行填补。思路不仅比原来简单，而且性能比原先提高了一倍。