Splitting the pandas dataset based on several condition

user_17103465 0
Viewed 40

I have two tables : Reference Table and Main Table (sample) as shown below :

Reference Table (ref)

MasterID    Myset
100      TRAIN
101      TRAIN
102      TRAIN
103      TRAIN
104      TRAIN
105      TRAIN
106      TEST
107      TEST
...

Main Table (df)

MasterID    MyId  Score Color 
100          123    34  Red
100          234    45  Pink
100          345    67  Blue
110          456    75  Green
115          567    36  Pink
115          678    78  Blue
104          679    35  Green
104          569    78  Pink
106          387    77  Blue
106          388    76  Pink
121          390    86  Red
122          450    94  Pink
123          333    102 Blue
123          222    111 Yellow
...

I want to add a new column "Myset" in the Main Table which I will be using to split the main table into two datasets : train and test.

# Splitting the dataset
train = df[df.Myset== 'TRAIN']
test = df[df.Myset== 'TEST']

I want to add the new column "Myset" based on the condition below:

  1. If a MasterID in the main table already belongs to a set in the reference table , the new column "Myset" will take the same value as the reference table.

  2. If a MasterID in the main table doesn't match with the reference table the new column "Myset" will either take 'TRAIN' or 'TEST' values , try to achieve an overall equal split (if possible) based on MyId which is a unique identifier.

  3. One MasterID can only belong to one set in the main table ; either 'TEST' or 'TRAIN'

Expected Output :

MasterID MyId Score Color   Myset
100      123    34    Red    TRAIN
100      234    45    Pink   TRAIN
100      345    67    Blue   TRAIN
110      456    75    Green  TEST
115      567    36    Pink   TRAIN
115      678    78    Blue   TRAIN
104      679    35    Green  TRAIN
104      569    78    Pink   TRAIN
106      387    77    Blue   TEST
106      388    76    Pink   TEST
121      390    86    Red    TEST
122      450    94    Pink   TEST
123      333    102   Blue   TEST
123      222    111   Yellow TEST
...
pandas python
1 Answers

You can merge the two dataframes using inner join on Master ID key:

import pandas as pd    
expected=pd.merge(ref,main,left_on='MasterID',right_on='MasterID')
user_59067750
Related