按列表列中的元素对 Pandas 数据框进行分组

Question

我试图通过按列表列中的元素分组来获取熊猫数据框中列的总和。我将创建一个虚拟数据集来表示我正在使用的数据：

preg_df = pd.DataFrame({'Diag_Codes': [['O1414', 'O4103X0', 'O365930', 'O76'], 
                                       ['O200', 'N3000', 'M545', 'R102', 'R110', 'Z3A01'],
                                       ['O365922', 'O30032', 'O09512', 'Z3A26'], 
                                       ['O2341', 'O200', 'Z3A01'], 
                                       ['O209', 'Z3A01']], 
                        'First_Trimester': [0, 1, 0, 1, 1], 
                        'Second_Trimester': [0, 0, 1, 0, 0], 
                        'Third_Trimester': [1, 0, 0, 0, 0]})

我想根据这些数据创建一个新数据框，该数据框按 preg_df 的“Diag_Codes”列中包含的诊断代码分组。我已经能够使用以下 for 循环来完成此操作：

# Create a list of unique diagnosis codes from the preg_df dataframe
diagnoses = list(set([item for sublist in preg_df.Diag_Codes.tolist() for item in sublist]))

diag_dfs = []

for i in diagnoses:
    
    diag_indices = []
    diag_df = pd.DataFrame()
    
    # Get the indices at which the diagnosis code exists within the 'Diag_Codes' column
    [diag_indices.append(index) for index, row in preg_df.iterrows() if i in preg_df.loc[index, 'Diag_Codes']]
    
    # Subset the dataframe to obtain only records in which the diagnosis code exists within 'Diag_Codes' column
    diag_df = preg_df.loc[diag_indices, 'First_Trimester':]
    diag_df['Diag_Code'] = i
    diag_df['Total_Cases'] = len(diag_indices)
    
    # Group by diagnosis code and the total number of cases and get the aggregate sum of all other columns
    diag_df = diag_df.groupby(['Diag_Code', 'Total_Cases']).sum()
    diag_dfs.append(diag_df)
    
diag_data = pd.concat(diag_dfs).sort_values(by=['Total_Cases'], ascending=False)
diag_data.head()

上面的 for 循环确实产生了我有兴趣从原始数据集创建的数据框，但是，这种方法不适用于大型数据集。我正在使用的实际数据框大约有 500 万行，并且包含数万个唯一的诊断代码。因此，我无法通过使用上面共享的 for 循环来获得我按诊断代码分组的所需数据帧。是否有更有效的方法可以让我使用更大的数据集来完成所需的输出？

Answer 1

让我们试试吧：

preg_df.explode('Diag_Codes').groupby('Diag_Codes').sum()

输出：

            First_Trimester  Second_Trimester  Third_Trimester
Diag_Codes                                                    
M545                      1                 0                0
N3000                     1                 0                0
O09512                    0                 1                0
O1414                     0                 0                1
O200                      2                 0                0
O209                      1                 0                0
O2341                     1                 0                0
O30032                    0                 1                0
O365922                   0                 1                0
O365930                   0                 0                1
O4103X0                   0                 0                1
O76                       0                 0                1
R102                      1                 0                0
R110                      1                 0                0
Z3A01                     3                 0                0
Z3A26                     0                 1                0

Answer 2

这应该可行。它会将每个列表 diag_code 与它们自己的列和行堆叠在一起，因此使用它们会更容易。

In:

diag_codes = pd.DataFrame(preg_df["Diag_Codes"].tolist()).stack()
diag_codes.index = diag_codes.index.droplevel(-1)
diag_codes.name = "diag_codes"
grouped_codes = preg_df.join(diag_codes).groupby('diag_codes').sum()
grouped_codes

输出：

Out:

    First_Trimester     Second_Trimester    Third_Trimester
diag_codes          
M545    1   0   0
N3000   1   0   0
O09512  0   1   0
O1414   0   0   1
O200    2   0   0
O209    1   0   0
O2341   1   0   0
O30032  0   1   0
O365922     0   1   0
O365930     0   0   1
O4103X0     0   0   1
O76     0   0   1
R102    1   0   0
R110    1   0   0
Z3A01   3   0   0
Z3A26   0   1   0

如有必要，请告诉我进行任何修复，或者您可以继续进行此操作。告诉我 te 数据库的维度，以便我查看它是否优化得很好。请记住，始终尝试使用内置函数并使用“for”循环作为最后的手段。