首页 > 解决方案 > numpy 或 pandas groupby 方式替换 2 个 for 循环

问题描述

我有以下问题有主要类的列表:

list_main_classes = [3,4]
data = pd.DataFrame({
    'label_col':[1,1,2,2,3,3,3,4,4], 
    'second_classes_column':[
        "class1", 
        "class2", 
        "class1", 
        "class2", 
        "class3", 
        "class3", 
        "class3", 
        "class4", 
        "class2"
    ]})

有一个列"second_classes_column" 我基本上想要做的是从列表"list_main_classes"中删除一些满足某些条件的元素。什么条件?

  1. 第二类列不能命名"certain_name"
  2. "second_classes_column"元素不得出现在由'label_col'. 这意味着对于由“label_col”的元素 4 创建的组,其中不得有"second_classes_column"出现在其他组中的元素。在我们的例子中,元素不满足这一点,"class2"因为它已经出现在之前(第 2 行和第 4 行)。因此,我们将从 list_main_classes 中删除 4,但保留 3,因为它满足所有要求,

问题是否有更快的方法来做到这一点,熊猫 groupby,numpy,已经用 2 个 for 循环完成了?

标签: pythonpandasnumpypandas-groupbynumpy-ndarray

解决方案


您应该对数据执行合并,然后对生成的数据框进行过滤。

此外,如果分配了"second_classes_column"多个唯一"label_col"值,则它是无效的,因此您可以预先计算与每个关联的 label_cols 的数量"second_classes_column"

# setup some useful variables
main_classes = pd.DataFrame({"main_classes": list_main_classes})
count_unique_classes = data.groupby("second_classes_column")["label_col"].nunique().to_dict()

def your_logic(x):
    second_id = x["second_classes_column"]
    label_col = x["label_col"]
    
    case1 = second_id != "certain_class"
    case2 = count_unique_classes[second_id] > 1
    
    return case1 and case2

# merge the two data frames
joint_df = pd.merge(data, main_classes, left_on="label_col", right_on="main_classes")

# now you can easily do the filter and perform your logic
to_drop = joint_df.apply(your_logic, axis=1)
list_main_indexes_to_drop = joint_df[to_drop].main_classes

所以结果:

>>> list_main_indexes_to_drop.values
... array([4])

可以使用 a filter、 set 操作或np.setdiff1d

>>> list(set(list_main_classes) - set(list_main_indexes_to_drop))
... [3]

或者

>>> np.setdiff1d(list_main_classes, list_main_indexes_to_drop)
... array([3])

更新。你可能不喜欢your_logicandapply所以你可以使用向量化的布尔运算来做到这一点,如下所示:

# setup some useful variables
main_classes = pd.DataFrame({"main_classes": list_main_classes})
count_unique_classes = data.groupby("second_classes_column")["label_col"].nunique().ge(2)
invalid_classes = set(count_unique_classes[count_unique_classes].index)

# merge the two data frames
joint_df = pd.merge(data, main_classes, left_on="label_col", right_on="main_classes")

# your logic
joint_df = joint_df[
    (joint_df.second_classes_column != "certain_class") & 
    (joint_df.second_classes_column.isin(invalid_classes)) 
]

# now you can easily do the filter and perform your logic
list_main_indexes_to_drop = joint_df.main_classes
list_main_indexes_to_drop.values

推荐阅读