msksantosh
msksantosh

Reputation: 399

All permutations and combinations of a pandas column based on the group it belongs to

I have a pandas dataframe which has city names and the state to which they belong. I am trying to obtain all the possible combinations of city name comparisons for each state.

Example dataframe:

City State
----------

LosA Cali 
SanJ Cali
SanF Cali
Char NC
Rale NC

Expected Answer:

City1 City2 State
----------
LosA  SanJ  Cali
LosA  SanF  Cali
SanJ  SanF  Cali
Char  Rale  NC

I have used combinations from itertools which gives the whole combinations, but is there a way to achieve based on the State as well?

Upvotes: 0

Views: 1077

Answers (1)

cs95
cs95

Reputation: 402813

Use a combination of groupby + itertools.combinations:

from itertools import combinations

g = df.groupby('State').apply(lambda x: 
      pd.Series(list(combinations(x.City, 2))))

df = pd.DataFrame(g.apply(list).tolist(), columns=['City1', 'City2'])
df['State'] = g.index.get_level_values(0)

df

  City1 City2 State
0  LosA  SanJ  Cali
1  LosA  SanF  Cali
2  SanJ  SanF  Cali
3  Char  Rale    NC

Upvotes: 3

Related Questions