replace loops with vector operations

Viewed 69

This is an example of a problem I need to solve. I have 2 tables And I have to find a substring(dff2) in a string(dff). And if substring in a string exists add this row to the output dataframe. The problem is that I wrote the code using loops and it works too slow for thousands of rows. How could I rewrite the code using pandas methods?

data = {'long string column': ['aaaabbbbccccdddd', 'bbbbccccddddeeee','ccccddddeeeeffff','ddddeeeeffffgggg'],}
dff = pd.DataFrame(data, columns = ['long string column'])


data2 = {'substring_column': ['aaaa', 'bbbb','cccc','dddd'], 'status': ['best', 'good','bad','worst'],}
dff2 = pd.DataFrame(data2, columns = ['substring_column','status'])


df_output = pd.DataFrame()
df_output_small = pd.DataFrame()


for i in range(len(dff2)):
    df_output_small = dff[(dff['long string column'].str.contains(dff2['substring_column'].iloc[i]))]
                          
    df_output_small['status'] = dff2['status'].iloc[i]

    df_output = df_output.append(df_output_small, ignore_index=True)
                          
df_output
1 Answers

Use Series.str.extractall for all substrings and then right merge:

pat = f'({"|".join(dff2["substring_column"])})'

If there are some substrings which should be regexes use re.escape:

import re
pat = '|'.join(re.escape(x) for x in dff2['substring_column'])

df = (dff.set_index('long string column', drop=False)['long string column']
         .str.extractall(pat)[0]
         .reset_index(level=1, drop=True)
         .reset_index(name='substring_column'))
df = df.merge(dff2, on='substring_column', how='right')
print (df)
  long string column substring_column status
0   aaaabbbbccccdddd             aaaa   best
1   aaaabbbbccccdddd             bbbb   good
2   bbbbccccddddeeee             bbbb   good
3   aaaabbbbccccdddd             cccc    bad
4   bbbbccccddddeeee             cccc    bad
5   ccccddddeeeeffff             cccc    bad
6   aaaabbbbccccdddd             dddd  worst
7   bbbbccccddddeeee             dddd  worst
8   ccccddddeeeeffff             dddd  worst
9   ddddeeeeffffgggg             dddd  worst

Last remove substring_column column:

df = df.drop('substring_column', axis=1)
Related