使用 DataFrame 和 Dictionary 作为输入编写函数来清理 python DataFrame

Question

假设我有一个 DataFrame，其中一列是“类别”，并且有 3 个类别：

1. cat_a 有 100 条记录
2. cat_b 有 5 条记录
3. cat_c 有 1000 条记录

我想编写一个函数来清理该 Dataframe。更具体地说，我想提供一个带有键、值的字典来告诉函数要清理哪些类别以及如何清理。假设我用这本字典提供函数：

categorysize = {"cat_b": 0,
                "cat_c": 100}

我的目标是：

1. 我想完全删除“cat_b”，这就是我设置 value=0 的原因
2. 我想将“cat_c”从 1000 条记录随机下采样到 100 条

我写了这个函数，但它不起作用。在所有情况下，我只设法删除一个类别或对另一个类别进行下采样。但不是两者都有:(

def df_reshaper(DF, categorysize):
    # I create a empty dataframe to write my data to                                                                                                                                               
    X = pd.DataFrame() 
    for key, value in categorysize.items():
        # If a key in "categorysize" has a value of 0 then drop that category from the data
        if value==0:
            DF = DF.loc[DF.category!=key, :]
            X = X.append(DF)
            return X
        # If a key in "categorysize" has a value > 0  then take a sample the size of value
        elif value>0:
            # First take a sample of x values for later
            DF2 = DF.loc[DF.category==key, : ].sample(value)
            # Then drop that category completely from the Dataframe
            DF3 = DF.loc[DF.category!=key]
            # Then add the sample from step x to the new dataframe
            DF4 = DF2.append(DF3)
            X = X.append(DF4)
            return X
    return X

生成的 DF 应该有 100 条 cat_a 记录和 100 条 cat_c 记录

Answer 1

问题在于您如何使用return，因为这将立即结束函数，即在您的if 和elif 条件结束后，仅执行其中一项操作后，您的函数将立即返回而无需执行其余操作的字典。在遍历字典并将列附加到此数据帧之前，您需要创建一个希望返回的空数据帧，该数据帧在循环外的函数结束时返回。