Given two dataframes with grouped items and another with equivalent groups, can I create a dataframe with all pairs of units from matching groups?

Question

Let's say I have a dataframe df1:

   df1_unit   base_unit
0     x         x
1     y         x
2     z         z
3     t         z
4     u         z

and another called df2:

   df2_unit   base_unit
0     a         b
1     b         b
2     c         c
3     d         e
4     e         e

and yet another dataframe df_eq that gives the equivalence between groups:

   df1_unit   df2_unit
0     x         b
1     z         e

In case of df1, the base_unit is basically the df1_unit that acts as a parent unit for all the units in a group. i.e. x is the base unit with which the x, y, z units are identified as a group. Similarly for df2.

I'm trying to generate a dataframe with all possible pairs of items from equivalent groups, but without the pairs in the df_eq dataframe (this is a restriction we impose). In this case, the unrestricted output would be:

   df1_unit   df2_unit
0     x         a
1     x         b       (shouldn't be included)
2     y         a
3     y         b
4     z         d
5     z         e       (shouldn't be included)
6     t         d
7     t         e
8     u         d
9     u         e

and the desired, restricted output would be:

   df1_unit   df2_unit
0     x         a
1     y         a
2     y         b
3     z         d
4     t         d
5     t         e
6     u         d
7     u         e

I'm having difficulty in generating even the unrestricted output without resorting to ridiculous brute force methods. Is there an efficient way to achieve the desired output?

EDIT: I've made some headway using the following code:

temp = dfeq.rename(columns={'df2u':'base_unit'}).merge(df2, on='base_unit', how='left')
temp = temp[['df1u', 'df2u']]
out = temp.rename(columns={'df1u':'base_unit'}).merge(df1, on='base_unit', how='left')
out = out[['df1u', 'df2u']]

Does this seem correct? Also, I'm not sure about how to remove the rows in out that are also present in dfeq.

user_4508767 · Answer

You can use the following approach:

We are going to map base_unit column in df1, taking into account the data in df_eq, for that we are going to use to_dict(orient='split') to build a dictionary which will be used by map method.
```
map_base_unit_df2 = dict(df_eq.to_dict(orient='split')['data']) 
df1['base_unit'] = df1['base_unit'].map(map_base_unit_df2)
```
We are going to build df_unrestricted using pd.merge selecting only the columns matter for us.
```
df_unrestricted = pd.merge(df1, df2, on='base_unit')[['df1_unit', 'df2_unit']]
```
Finally, adding the final restriction, I mean, we are going to drop the records present in df_eq, using pd.merge + drop_duplicates(keep=False) + reset_index(drop=True)
```
df_output = pd.concat([df_unrestricted, df_eq]).drop_duplicates(keep=False).reset_index(drop=True)
```

Full code:

map_base_unit_df2 = dict(df_eq.to_dict(orient='split')['data']) 
df1['base_unit'] = df1['base_unit'].map(map_base_unit_df2)
df_unrestricted = pd.merge(df1, df2, on='base_unit')[['df1_unit', 'df2_unit']]
df_output = pd.concat([df_unrestricted, df_eq]).drop_duplicates(keep=False).reset_index(drop=True)
print(df_output)

Output:

   df1_unit df2_unit
0      x       a
1      y       a
2      y       b
3      z       d
4      t       d
5      t       e
6      u       d
7      u       e

Given two dataframes with grouped items and another with equivalent groups, can I create a dataframe with all pairs of units from matching groups?

1 Answers