Search code examples
pythoncombinationspermutationpython-itertools

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?


Solution

  • 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