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

user_13204346 0
Viewed 260

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.

pandas-groupby pandas dataframe python
2 Answers

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
user_42384080

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)

user_63615310
Related