Is there an optimal way to get all combinations of values in a grouped pandas dataframe?

Question

Sample dataframe as below

df = pd.DataFrame({'ID': ['a', 'a', 'a', 'b', 'b', 'c', 'c'], 
                   'color': ['red', 'blue', 'green', 'red', 'blue', 'red', 'green']})

I want 2 columns with all combinations of the color field after grouping by ID.
I want the resultant dataframe as shown below

ID	color1	color2
a	red	blue
a	red	green
a	blue	red
a	blue	green
a	green	red
a	green	blue
b	red	blue
b	blue	red
c	red	green
c	green	red

I have tried using itertools.permutations but am looking for something more direct or for a solution that utilizes Pandas more.

Quang Hoang · Accepted Answer · 2021-01-21 19:20:40Z

7

I think you can do a self merge and query:

df.merge(df, on='ID', suffixes=[1,2]).query('color1 != color2')

Or similar, merge then filter:

(df.merge(df, on='ID', suffixes=[1,2])
   .loc[lambda x: x['color1'] != x['color2']]
)

Output:

   ID color1 color2
1   a    red   blue
2   a    red  green
3   a   blue    red
5   a   blue  green
6   a  green    red
7   a  green   blue
10  b    red   blue
11  b   blue    red
14  c    red  green
15  c  green    red

answered Jan 21, 2021 at 19:20

Quang Hoang

151k11 gold badges64 silver badges86 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Scott Boston · Accepted Answer · 2021-01-21 20:30:59Z

2

You can use this method:

from itertools import permutations

s = df.groupby('ID')['color']\
      .apply(lambda x: list(permutations(x, 2))).explode()
dfi= pd.DataFrame().from_records(s, index=s.index, columns=['color1', 'color2'])
dfi

Output:

   color1 color2
ID              
a     red   blue
a     red  green
a    blue    red
a    blue  green
a   green    red
a   green   blue
b     red   blue
b    blue    red
c     red  green
c   green    red

Timings:

#This method

3.18 ms ± 23.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

#Self join method

5.96 ms ± 105 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

answered Jan 21, 2021 at 20:30

Scott Boston

154k15 gold badges160 silver badges208 bronze badges

Collectives™ on Stack Overflow

Is there an optimal way to get all combinations of values in a grouped pandas dataframe?

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related