Pandas: set difference by group

user_1393043 0
Viewed 225

I have a very large dataset containing the members in each team in each month. I want to find additions and deletions to each team. Because my dataset is very big, I'm trying to use in-built functions as much as possible.

My dataset looks like this:

  month team    members
0   0   A   X, Y, Z
1   1   A   X, Y
2   2   A   W, X, Y
3   0   B   D, E
4   1   B   D, E, F
5   2   B   F

It's generated by the following code:

num_months = 3
num_teams = 2
obs = num_months*num_teams

df = pd.DataFrame({"month": [i % num_months for i in range(obs)],
                  "team": ['AB'[i // num_months] for i in range(obs)],
                   "members": ["X, Y, Z", "X, Y", "W, X, Y", "D, E", "D, E, F", "F"]})
df

The result should be like this:

    month   team    members additions   deletions
0   0       A       X, Y, Z None    None
1   1       A       X, Y    None    Z
2   2       A       W, X, Y W       None
3   0       B       D, E    None    None
4   1       B       D, E, F F       None
5   2       B       F       None    D, E

or in Python code

df = pd.DataFrame({"month": [i % num_months for i in range(obs)],
                  "team": ['AB'[i // num_months] for i in range(obs)],
                   "members": ["X, Y, Z", "X, Y", "W, X, Y", "D, E", "D, E, F", "F"],
                  "additions": [None, None, "W", None, "F", None],
                   "deletions": [None, "Z", None, None, None, "D, E"]
                  })

A technique that immediately comes to mind is to create a new column which shows the lagged value of members in each group, followed by taking the set difference (both ways) between both columns.

Is there a way to take set differences between columns using pandas inbuilt functions?

Are there other techniques I should try?

pandas python
2 Answers

Using set, groupby, apply, and shift.

  • For efficiency:
    • Convert members to set type because - is an unsupported operand, which will cause a TypeError.
    • Leave additions and deletions as set type

Using apply

  • With a dataframe of 60000 rows:
    • 91.4 ms ± 2.77 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
# clean the members column
df.members = df.members.str.replace(' ', '').str.split(',').map(set)

# create del and add
df['deletions'] = df.groupby('team')['members'].apply(lambda x: x.shift() - x)
df['additions'] = df.groupby('team')['members'].apply(lambda x: x - x.shift())

# result
 month team    members additions deletions
     0    A  {Z, X, Y}       NaN       NaN
     1    A     {X, Y}        {}       {Z}
     2    A  {W, X, Y}       {W}        {}
     0    B     {D, E}       NaN       NaN
     1    B  {D, F, E}       {F}        {}
     2    B        {F}        {}    {D, E}

More Efficiently

  • pandas.DataFrame.diff
  • With a dataframe of 60000 rows:
    • 60.7 ms ± 3.54 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
df['deletions'] = df.groupby('team')['members'].diff(periods=-1).shift()
df['additions'] = df.groupby('team')['members'].diff()
user_77588040

Here is one way to do it. Not sure if this is the most efficient. I've found that is not that straightforward to optimize pandas performance by just looking at the code.

The strategy I've adopted is to calculate the deletions and additions separately and then somehow merge that information back into the original DataFrame.

This solution assumes that the input DataFrame is sorted by (team, month). If not, you'd need to do that first.

def set_diff_adds(x):
  retval = {}
  for m, b, a in zip(x.month.iloc[1:], x.members.iloc[1:], x.members):
    retval[m] = (set(b.replace(' ', '').split(',')) - 
                 set(a.replace(' ', '').split(',')))
  return retval

def set_diff_dels(x):
  retval = {}
  for m, b, a in zip(x.month.iloc[1:], x.members.iloc[1:], x.members):
    retval[m] = (set(a.replace(' ', '').split(',')) - 
                 set(b.replace(' ', '').split(',')))
  return retval

deletions = df.groupby('team').apply(set_diff_dels).apply(pd.Series)
deletions.columns.set_names('month', inplace=True)
deletions = deletions.stack().to_frame('deletions').reset_index()

merged = df.merge(deletions, how='outer')

additions = df.groupby('team').apply(set_diff_adds).apply(pd.Series)
additions.columns.set_names('month', inplace=True)
additions = additions.stack().to_frame('additions').reset_index()

merged = merged.merge(additions, how='outer')

merged


   month team  members deletions additions
0      0    A  X, Y, Z       NaN       NaN
1      1    A     X, Y       {Z}        {}
2      2    A  W, X, Y        {}       {W}
3      0    B     D, E       NaN       NaN
4      1    B  D, E, F        {}       {F}
5      2    B        F    {D, E}        {}
user_19175000
Related