Compare a Column value to different columns and return a value from same row but different column using Python Pandas

Viewed 92

I got data as in image 1.
enter image description here

By using for loop in Python Pandas, I have splitted FG and PID columns as in image 2.
enter image description here

Based on the condition, I want the output as picture 3:

if Main_FG == FG_1, value will be from PID_1
if Main_FG == FG_2, value will be from PID_2
if Main_FG == FG_3, value will be from PID_3

enter image description here

I have tried using for loop but does not help.

What would be the effective way to do that in Python? Please note that, the number of FG_1, FG_xxxxx, varies therefore, I need code to be dynamic and I am unable to do by using np.select().

6 Answers

My approach is to define an external function to select the relevant PID:

import pandas as pd

df = pd.DataFrame({'Main_FG': ['AB', 'AC', 'AM', 'AP', 'AE', 'RE', 'Rw', 'NA'],
                   'FG': ['CD,AB,QF', 'AC', 'AM', 'AD, GH, AP', 'GH, ae, ME', 'RE', 'RW', 'RW'],
                   'PID': ['A.900,B1.56,C1.650', 'R1', 'Q1.45', 'S1.76,S2.56,S3', 'J1.33,J2.39,S3', 'B1.09, B2.56, E2', 'S3.07', 'S3.07']  })
print(df)

  Main_FG          FG                 PID
0      AB    CD,AB,QF  A.900,B1.56,C1.650
1      AC          AC                  R1
2      AM          AM               Q1.45
3      AP  AD, GH, AP      S1.76,S2.56,S3    #note the whitespaces between commas
4      AE  GH, ae, ME      J1.33,J2.39,S3    #note the lowercase
5      RE          RE    B1.09, B2.56, E2    #note the whitespaces between commas
6      Rw          RW               S3.07    #note the lowercase
7      NA          RW               S3.07    #not found

The function also takes care of 3 data cleaning scenarios:
(1) to be case insensitive by matching based on lowercase,
(2) to remove whitespaces when splitting on comma,
(3) to return None when match is not found (thanks @CyrillePontvieux!)

def select_pid(row):
    try:
        idx = [e.strip().lower() for e in row['FG'].split(',')].index(row['Main_FG'].lower())
        return [e.strip() for e in row['PID'].split(',')][idx]
    except:
        return None
    
df['SELECTED_PID'] = df.apply(lambda row: select_pid(row), axis=1)
print(df)

  Main_FG          FG                 PID SELECTED_PID
0      AB    CD,AB,QF  A.900,B1.56,C1.650        B1.56
1      AC          AC                  R1           R1
2      AM          AM               Q1.45        Q1.45
3      AP  AD, GH, AP      S1.76,S2.56,S3           S3
4      AE  GH, ae, ME      J1.33,J2.39,S3        J2.39
5      RE          RE    B1.09, B2.56, E2        B1.09
6      Rw          RW               S3.07        S3.07
7      NA          RW               S3.07         None

You can use something like that:

from pandas import DataFrame, concat

df=DataFrame(data={'Main_FG': ['AB', 'AC', 'AM', 'AP', 'AE', 'RE', 'Rw'], 'FG_1': ['CD', 'AC', 'AM', 'AD', 'GH', 'RE', 'RW'], 'FG_2': ['AB', None, None, 'GH', 'AE', None, None], 'FG_3': ['QF', None, None, 'AP', 'ME', None, None], 'PID_1': ['A.900', 'R1', 'Q1.45', 'S1.76', 'J1.33', 'B1.09', 'S3.07'], 'PID_2': ['B1.56', None, None, 'S2.56', 'J2.39', 'B2.56', None], 'PID_3': ['C1.650', None, None, 'S3', 'S3', 'E2', None]})
print(df)

  Main_FG FG_1  FG_2  FG_3  PID_1  PID_2   PID_3
0      AB   CD    AB    QF  A.900  B1.56  C1.650
1      AC   AC  None  None     R1   None    None
2      AM   AM  None  None  Q1.45   None    None
3      AP   AD    GH    AP  S1.76  S2.56      S3
4      AE   GH    AE    ME  J1.33  J2.39      S3
5      RE   RE  None  None  B1.09  B2.56      E2
6      Rw   RW  None  None  S3.07   None    None
for col in df.columns:
    df[col] = df[col].str.upper()
fg_cols = [col for col in df.columns if col.startswith('FG_')]
pid_cols = [col for col in df.columns if col.startswith('PID_')]
assert len(fg_cols) == len(pid_cols)
df_list = []
for i, fg_col in enumerate(fg_cols):
    df_list.append(df[df['Main_FG'] == df[fg_col]][['Main_FG', pid_cols[i]]].rename(columns={pid_cols[i]: 'SELECTED_PID'}))
result_df = concat(df_list).sort_index()
print(result_df)
  Main_FG SELECTED_PID
0      AB        B1.56
1      AC           R1
2      AM        Q1.45
3      AP           S3
4      AE        J2.39
5      RE        B1.09
6      RW        S3.07

You could sort the Main_FG column if you like.

Another option would be to use melt and assign dummy columns based on FG_i

df_ = (df
    .melt(id_vars=['Main_FG', 'FG_1', 'FG_2', 'FG_3'])
    .assign(
        FG_1_ = lambda ser_: np.where(ser_.Main_FG == ser_.FG_1, ser_.variable == 'PID_1', np.nan),
        FG_2_ = lambda ser_: np.where(ser_.Main_FG == ser_.FG_2, ser_.variable == 'PID_2', np.nan),
        FG_3_ = lambda ser_: np.where(ser_.Main_FG == ser_.FG_3, ser_.variable == 'PID_3', np.nan)
    )
    .query("FG_1_ == 1 | FG_2_ == 1 | FG_3_ == 1")
    .rename(columns={'value': 'SELECTED_PID'})
    .loc[:, ['Main_FG', 'SELECTED_PID']]
)

   Main_FG SELECTED_PID
1       AC           R1
2       AM        Q1.45
5       RE        B1.09
6       RW        S3.07
7       AB        B1.56
10      AP        S2.56
11      AE        J2.39

Using a dummy example as yours is not reproducible.

You can use a list comprehension to loop over the key/values as a dictionary:

df['selected_PID'] = [dict(zip(k,v)).get(ref)
                      for ref,k,v in zip(df['main_FG'].str.upper(),
                                         df['FG'].str.split(','),
                                         df['PID'].str.split(','))]

NB. If you want to drop the columns, use df.pop('FG').str.split(',') in place of df['FG'].str.split(',') (same for PID).

Alternative with explode:

(df.assign(FG=df['FG'].str.split(','),
          PID=df['PID'].str.split(',')
         )
   .explode(['FG', 'PID'])
   .loc[lambda x: x['main_FG'].eq(x['FG'])]
   .rename(columns={'PID': 'selected_PID'})
   .drop(columns='FG')
)

output:

  main_FG     FG    PID selected_PID
0       B  A,B,C  U,V,W            V
1       E    D,E    X,Y            Y
2       F      F      Z            Z

used input:

  main_FG     FG    PID
0       B  A,B,C  U,V,W
1       E    D,E    X,Y
2       F      F      Z

I am going from situation 2 to situation 3:

You can split your dataframe into two subdataframes. df1 containing FG_... and df2 containing PID_... columns. Create a mask using Main_FG and df1 and apply this on df2.

df = pd.DataFrame([
    ['AB', 'CD', 'AB', 'QF', 'A.900', 'B1.56', 'C1.650'],
    ['AC', 'AC', '', '', 'R1', '', ''],
    ['AM', 'AM', '', '', 'Q1.45', '', ''],
    ['AP', 'AD', 'GH', 'AP' , 'S1.76', 'S2.56', 'S3']],
    columns=['Main_FG', 'FG_1', 'FG_2', 'FG_3', 'PID_1', 'PID_2', 'PID_3'])

n = len(df) // 2 + 1
df1 = df.iloc[:, 1:n]
df2 = df.iloc[:, n:]

mask = df1.eq(df.Main_FG, axis='rows')
mask[mask == 0] = np.nan

df['SELECTED_PID'] = (df2.values * mask).bfill(axis=1).iloc[:, 0]
df3 = df[['Main_FG', 'SELECTED_PID']]

Output df3:

  Main_FG SELECTED_PID
0      AB        B1.56
1      AC           R1
2      AM        Q1.45
3      AP           S3

It looks like problem described may be solved with a sequence of str.split and explode method

df['FG'] = df['FG'].str.split(',')
df['PID'] = df['PID'].str.split(',')
# Workaround to Assume Different Lists' Lengths
for _ in range(df.shape[0]):
    if len(df.iloc[_, 1]) < len(df.iloc[_, 2]):
        df.iloc[_, 1] = len(df.iloc[_, 2]) * [df.iloc[_, 1]]
df = df.explode(['FG', 'PID'])
print(df)

Thus, you may arrive at the flattened DataFrame

  Main_FG  FG     PID
0      AB  CD   A.900
0      AB  AB   B1.56
0      AB  QF  C1.650
1      AC  AC      R1
2      AM  AM   Q1.45
3      AP  AD   S1.76
3      AP  GH   S2.56
3      AP  AP      S3
4      AE  GH   J1.33
4      AE  AE   J2.39
4      AE  ME      S3
5      RE  RE   B1.09
5      RE  RE   B2.56
5      RE  RE      E2
6      Rw  RW   S3.07

doing everything you want from there.

Related