在组内联合非集迭代的有效方法

Question

I have this df

df = pd.DataFrame(dict(
        A=['b', 'a', 'b', 'c', 'a', 'c', 'a', 'c', 'a', 'a'],
        B=[[0, 2, 3, 1],
           [9, 6, 7, 2],
           [6, 0, 1, 4],
           [9, 2, 5, 1],
           [5, 1, 4, 8],
           [8, 5, 6, 6],
           [0, 9, 0, 0],
           [2, 6, 1, 8],
           [7, 3, 2, 6],
           [8, 7, 1, 9]]
        ))

I want to group by 'A' and union all the lists in 'B'

Neither df.groupby('A').B.union() nor df.groupby('A').B.apply(set.union) work.

I want the result to be

A
a    {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
b                {0, 1, 2, 3, 4, 6}
c                {1, 2, 5, 6, 8, 9}
Name: B, dtype: object

Answer 1

The problem is that you need to cast them as sets first before applying the union. One solution would be to use sum to concatenate the groups, then cast to set using map

In [28]: df.groupby('A').B.sum().map(set)
Out[28]:
A
a    {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
b                {0, 1, 2, 3, 4, 6}
c                {1, 2, 5, 6, 8, 9}
dtype: object

Answer 2

maxymoo's answer is nice, but since it first adds all the lists together it might unnecessarily take a lot of memory (especially so if there are lots of duplicates).

Instead, you should first convert column B to sets, after which you can reduce to a single set much more efficiently. Like this:

df['B'] = df['B'].map(set)

   A             B
0  b  {0, 1, 2, 3}
1  a  {9, 2, 6, 7}
2  b  {0, 1, 4, 6}
3  c  {9, 2, 5, 1}
4  a  {8, 1, 4, 5}
5  c     {8, 5, 6}
6  a        {0, 9}
7  c  {8, 1, 2, 6}
8  a  {2, 3, 6, 7}
9  a  {8, 1, 9, 7}

df.groupby('A').B.apply(lambda x: reduce(set.union, x))

A
a    {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
b                {0, 1, 2, 3, 4, 6}
c                {1, 2, 5, 6, 8, 9}
Name: B, dtype: object

Or, as a one-liner, as maxymoo points out:

df.groupby('A').B.apply(lambda x : reduce(set.union, x.map(set)))

Answer 3

I'd use a function to apply with

def f(x):
    # grabbing first one so I can
    # make a set out of it
    first, *rest = x.values.tolist()
    # union won't work unless it's on
    # a set, it doesn't care about the rest
    return set(first).union(*rest)

df.groupby('A').B.apply(f)

A
a    {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
b                {0, 1, 2, 3, 4, 6}
c                {1, 2, 5, 6, 8, 9}
Name: B, dtype: object