Improve efficiency of iloc - merging two tables with startwith()

Viewed 157

I am trying to merge two tables such as: df1:

ID
A1A1A1
A1A1A2
B1B1B1

df2:

ID Country
A1A1A1 France
B Egypt
C1C Egypt

In Table 2, when the ID = B, it means that all IDs starting with B will have the same country. The same applies to ID = CCC1 Hence, I cannot use pd.merge as I cannot merge on an exact match. I have written a code that seems to work (using debug) but that is incredibly slow. I am looking therefore for faster solution. My df1 has ~80K rows, df2 has ~7K rows.

Expected Output: df2:

ID Country
A1A1A1 France
A1A1A2 nan
B1B1B1 Egypt

Here is what I did:

for i in range(len(df2)):
   for j in range(len(df1)):
      if df1['ID'].iloc[j].startswith(df2['ID'].iloc[i]):
        df1['Country'].iloc[j] = df2['Country'].iloc[i]

Thanks!

1 Answers

One solution to your problem is to build temporary columns with all (possible) different lengths. In your case 6 columns. Then you can convert you df2 to a dictionary and just look the ids up in the dictionary. Then you combine the columns with combine_first. Be aware that the order of your column list matters with combine_first.

import pandas as pd

ids = ['A1A1A1', 'A1A1A2', 'B1B1B1']
df1 = pd.DataFrame(ids, columns=['ID'], index=range(3))
df2 = pd.DataFrame.from_dict({'ID': {0: 'A1A1A1', 1: 'B', 2: 'C1C'},
                              'Country': {0: 'France', 1: 'Egypt', 2: 'Egypt'}})

# build dictionary from df2 (dictionary is probably faster than .loc). Also it is cleaner
map_id_dict = df2.set_index('ID')['Country'].to_dict()

# Define target column
df1['Country'] = None
# Build temporary columns
cols = [f'ID_{i}' for i in range(1, 7)]
for i, col in enumerate(cols):
    # lookup ids in dictionary from df2
    df1[col] = df1['ID'].str[:i + 1].apply(lambda x: map_id_dict.get(x))
    df1['Country'] = df1['Country'].combine_first(df1[col])
# drop temporary columns
df1 = df1.drop(columns=cols)

Output:

       ID Country
0  A1A1A1  France
1  A1A1A2    None
2  B1B1B1   Egypt
Related