根据一个列值在另一列中包含多个类别值拆分 dataframe

Question

I am trying to create two dataframes from the following data:

df = pd.DataFrame({'Product':['Prod1','Prod2','Prod3','Prod2','Prod5','Prod3']*4, 
                  'Inv_Type': ['X', 'Y']*12,
                 'Quant': np.random.randint(2,20, size=24)})

df.sort_values('Product', inplace=True, ignore_index=True) --Help with visual

They need to be separated based on whether the Products have both an X and Y associated with them or just all X's or all Y's.

Desired Output:

df1 = df[df['Product'] == 'Prod3']
df2 = df[df['Product'].str.contains('Prod1|Prod2|Prod5', na=False)]

I have tried numerous groupby attempts with filters, but I am obviously missing something.

Answer 1

m = df.groupby("Product")["Inv_Type"].transform(lambda x: len(x.unique()) == 1)

df1 = df[~m]
df2 = df[m]
print(df1)
print(df2)

Prints:

   Product Inv_Type  Quant
12   Prod3        X      4
13   Prod3        Y     18
14   Prod3        Y     11
15   Prod3        X      5
16   Prod3        Y      5
17   Prod3        X      3
18   Prod3        X     16
19   Prod3        Y     11

   Product Inv_Type  Quant
0    Prod1        X      5
1    Prod1        X      6
2    Prod1        X      8
3    Prod1        X     17
4    Prod2        Y      3
5    Prod2        Y     13
6    Prod2        Y      9
7    Prod2        Y      8
8    Prod2        Y      7
9    Prod2        Y      5
10   Prod2        Y     18
11   Prod2        Y     11
20   Prod5        X      4
21   Prod5        X     15
22   Prod5        X     10
23   Prod5        X      6

Answer 2

you can create a custom boolean to groupby and create two separate data frames inside a dictionary. Assuming that there are only two values in your Inv_Type so we can use nunique to fidn any group that has more than one value.

dfs = {int(grp) : data for grp,data 
          in df.groupby([df.groupby('Product')['Inv_Type'].transform('nunique') > 1])}


print(dfs[1])

   Product Inv_Type  Quant
12   Prod3        X      2
13   Prod3        Y     12
14   Prod3        Y      2
15   Prod3        X     19
16   Prod3        Y      6
17   Prod3        X      5
18   Prod3        X      4
19   Prod3        Y     13

print(dfs[0])

   Product Inv_Type  Quant
0    Prod1        X     16
1    Prod1        X     13
2    Prod1        X      8
3    Prod1        X     16
4    Prod2        Y     14
5    Prod2        Y     10
6    Prod2        Y      4
7    Prod2        Y     13
8    Prod2        Y      7
9    Prod2        Y     16
10   Prod2        Y     13
11   Prod2        Y     11
20   Prod5        X     11
21   Prod5        X     10
22   Prod5        X     13
23   Prod5        X     10

Answer 3

We can also do it with boolean mask and Pandas built-in aggregate function (for better execution speed) instead of custom lambda function (which is not optimized and slow), as follows:

mask = df.groupby("Product")["Inv_Type"].transform('nunique') > 1
df1 = df[mask]
df2 = df[~mask]

Result:

print(df1)


   Product Inv_Type  Quant
12   Prod3        X     15
13   Prod3        Y     19
14   Prod3        Y     16
15   Prod3        X     12
16   Prod3        Y      9
17   Prod3        X      8
18   Prod3        X      8
19   Prod3        Y      7



print(df2)


   Product Inv_Type  Quant
0    Prod1        X     17
1    Prod1        X     12
2    Prod1        X      9
3    Prod1        X      9
4    Prod2        Y      2
5    Prod2        Y     16
6    Prod2        Y     16
7    Prod2        Y      9
8    Prod2        Y     17
9    Prod2        Y     12
10   Prod2        Y     12
11   Prod2        Y     13
20   Prod5        X      2
21   Prod5        X     19
22   Prod5        X     16
23   Prod5        X     18