代码之家  ›  专栏  ›  技术社区  ›  Giora Simchoni

带集合交集的GroupBy和Aggregate

  •  2
  • Giora Simchoni  · 技术社区  · 7 年前

    我有一个带集合列的熊猫数据框:

    import pandas as pd
    
    df = pd.DataFrame({'group_var': [1,1,2,2], 'sets_var': [set([0, 1]), set([1, 2]), set([3, 4]), set([5, 6, 7])]})
    df
    
       group_var sets_var
    0          1      {0, 1}
    1          1      {1, 2}
    2          2      {3, 4}
    3          2   {5, 6, 7}
    

    我希望 groupby 这个 group_var 得到所有对应集合的交集 sets_var ,像这样:

       group_var sets_var
    0          1      {1}
    1          2      {}
    

    或者类似的系列:

       sets_var
    1  {1}
    2  {}
    

    我该如何优雅地进行呢?性能是第一要务。

    1 回复  |  直到 7 年前
        1
  •  4
  •   cs95 abhishek58g    7 年前

    使用 groupby , agg 减少使用 set.intersection .

    df.groupby('group_var', as_index=False).agg(lambda x: set.intersection(*x))
    
       group_var sets_var
    0          1      {1}
    1          2       {}
    

    如果性能是绝对重要的,我们可以尝试摆脱 lambda :

    from functools import partial, reduce 
    import operator
    
    p = partial(reduce, operator.and_)
    df.groupby('group_var', as_index=False).agg(p)
    
       group_var sets_var
    0          1      {1}
    1          2       {}
    

    但是,这只会执行成对交叉,因此您的里程可能会有所不同。


    或者,作为一个系列,

    pd.Series({
        k: set.intersection(*g.tolist()) 
        for k, g in df.groupby('group_var')['sets_var']})
    
    1    {1}
    2     {}
    dtype: object