Pandas vectorized way to produce pair-wise combinations of a group rows with all other rows

Viewed 66

Is there a pandas vectorized way to produce a certain subset of all pair-wise rows combinations as follows: given distinguished group of rows, I want to pair each row from the group with all other rows (i.e. with both ex-group and inside group rows). Thus if the whole set is of length n and the group of length k, k << n then I'm looking for a vectorized O(nk) solution.

For example, suppose we are given the following data frame

   CarMaker Model  HorsePower  TopSpeed
0      Audi    S3         100       200
1      Audi    S5         110       210
2       BMW    M3         120       220
3       BMW    M4         130       230
4  Mercedes   GLS         140       240
5  Mercedes   AMG         150       250

from copy-friendly piece of code

input_df = pd.DataFrame({
    "CarMaker": ["Audi", "Audi", "BMW", "BMW", "Mercedes", "Mercedes" ],
    "Model": ["S3", "S5", "M3", "M4", "GLS", "AMG"],
    "HorsePower": [100, 110, 120, 130, 140, 150],
    "TopSpeed": [200, 210, 220, 230, 240, 250]
})

and distinguished group being Audi cars, I want to pair all Audi models with all other rows to get

  CarMaker_main Model_main CarMaker_other Model_other  HP_main  HP_other  TopSpeed_main  TopSpeed_other

0          Audi         S3           Audi          S5      100       110            200             210
1          Audi         S3            BMW          M3      100       120            200             220
2          Audi         S3            BMW          M4      100       130            200             230
3          Audi         S3       Mercedes         GLS      100       140            200             240
4          Audi         S3       Mercedes         AMG      100       150            200             250
5          Audi         S5            BMW          M3      110       120            210             220
6          Audi         S5            BMW          M4      110       130            210             230
7          Audi         S5       Mercedes         GLS      110       140            210             240
8          Audi         S5       Mercedes         AMG      110       150            210             250
1 Answers

If you don't mind comparing the model to itself then you could use merge.

import pandas as pd

input_df = pd.DataFrame({
    "CarMaker": ["Audi", "Audi", "BMW", "BMW", "Mercedes", "Mercedes" ],
    "Model": ["S3", "S5", "M3", "M4", "GLS", "AMG"],
    "HorsePower": [100, 110, 120, 130, 140, 150],
    "TopSpeed": [200, 210, 220, 230, 240, 250]
})

df_main = input_df[input_df["CarMaker"]=="Audi"].copy()
df_other = input_df.copy()

df_main["key"] = 0
df_other["key"] = 0

pd.merge(df_main, df_other, how="outer", on="key", suffixes=("_main", "_other")).drop("key", axis=1).sort_index(axis=1)

If you have pandas > 1.2 you can skip the "key" shenanigans and just pass how="cross" to pandas merge

pd.merge(df_main, df_other, how="cross", suffixes=("_main", "_other")).sort_index(axis=1)
Related