Combined group by using pandas

Viewed 129

Imagine a pandasdata frame given by

df = pd.DataFrame({
    'id': range(1, 10),
    'mfr': ('a', 'b', 'a', 'c', 'd', 'e', 'd', 'd', 'f'),
    'vmn': ('A', 'A', 'B', 'C', 'D', 'E', 'F', 'F', 'D')
})

which gives the following table

   id mfr vmn
0   1   a   A
1   2   b   A
2   3   a   B
3   4   c   C
4   5   d   D
5   6   e   E
6   7   d   F
7   8   d   F
8   9   f   D

I wish to determine which id's belong to eachother by grouping either by mfrand/or vmn. I can easily assign a group id by using one of the other by

df['groupby_mfr'] = df.groupby('mfr').grouper.group_info[0]
df['groupby_vmn'] = df.groupby('vmn').grouper.group_info[0]

which gives the following

   id mfr vmn  groupby_mfr  groupby_vmn
0   1   a   A            0            0
1   2   b   A            1            0
2   3   a   B            0            1
3   4   c   C            2            2
4   5   d   D            3            3
5   6   e   E            4            4
6   7   d   F            3            5
7   8   d   F            3            5
8   9   f   D            5            3

Now I want to combine this to a new group id so the resulting data frame becomes like this

   id mfr vmn  groupby_mfr  groupby_vmn  combined_group
0   1   a   A            0            0               0
1   2   b   A            1            0               0
2   3   a   B            0            1               0
3   4   c   C            2            2               1
4   5   d   D            3            3               2
5   6   e   E            4            4               3
6   7   d   F            3            5               2
7   8   d   F            3            5               2
8   9   f   D            5            3               2

The first two rows are the same since vmn are equal. The third are also the same group since row 3 and 1 are the same for vmn. And so on...

Note also that this will be run on multiple columns with many rows so performance is much appreciated as well.

1 Answers

As suggested in the comments in the original post it can be solved by using networkx.

import networkx as nx
import pandas as pd

df = pd.DataFrame({
    'id': range(1, 10),
    'mfr': ('a', 'b', 'a', 'c', 'd', 'e', 'd', 'd', 'f'),
    'vmn': ('A', 'A', 'B', 'C', 'D', 'E', 'F', 'F', 'D')
})

G = nx.from_pandas_edgelist(df, 'mfr', 'vmn')
Gcc = nx.connected_components(G)

connected_map = dict()
for g, ids in enumerate(Gcc):
    for id in ids:
        connected_map[id] = g

df['combined_group'] = df['mfr'].map(connected_map)

which yields

   id mfr vmn  combined_group
0   1   a   A               0
1   2   b   A               0
2   3   a   B               0
3   4   c   C               1
4   5   d   D               2
5   6   e   E               3
6   7   d   F               2
7   8   d   F               2
8   9   f   D               2
Related