Pandas change values in column based on values in other column

user_8921867 0
Viewed 150

I have a dataframe in which one column represents some data, the other column represents indices on which I want to delete from my data. So starting from this:

import pandas as pd
import numpy as np

df = pd.DataFrame({'data':[np.arange(1,5),np.arange(3)],'to_delete': [np.array([2]),np.array([0,2])]})
df
>>>> data       to_delete
     [1,2,3,4]    [2]
     [0,1,2]     [0,2]

This is what I want to end up with:

new_df
>>>>   data     to_delete
     [1,2,4]       [2]
       [1]        [0,2]

I could iterate over the rows by hand and calculate the new data for each one like this:

new_data = []
for _,v in df.iterrows():
    foo = np.delete(v['data'],v['to_delete'])
    new_data.append(foo)
df.assign(data=new_data)

but I'm looking for a better way to do this.

pandas python
3 Answers

The overhead from calling a numpy function for each row will really worsen the performance here. I'd suggest you to go with lists instead:

df['data'] = [[j for ix, j in enumerate(i[0]) if ix not in i[1]] 
              for i in df.values]

print(df)

       data to_delete
0  [1, 2, 4]       [2]
1        [1]    [0, 2]

Timings on a 20K row dataframe:

df_large = pd.concat([df]*10000, axis=0)

%timeit [[j for ix, j in enumerate(i[0]) if ix not in i[1]] 
         for i in df_large.values]
# 184 ms ± 12.4 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit 
new_data = []
for _,v in df_large.iterrows():
    foo = np.delete(v['data'],v['to_delete'])
    new_data.append(foo)

# 5.44 s ± 233 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit df_large.apply(lambda row: np.delete(row["data"], 
                       row["to_delete"]), axis=1)
# 5.29 s ± 340 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
user_96986840

You should use the apply function in order to apply a function to every row in the dataframe:

df["data"] = df.apply(lambda row: np.delete(row["data"], row["to_delete"]), axis=1)
user_48522900

An other solution based on starmap:

This solution is based on a less known tool from the itertools module called starmap.

Check its doc, it's worth a try! 

import pandas as pd
import numpy as np
from itertools import starmap

df = pd.DataFrame({'data': [np.arange(1,5),np.arange(3)],
                   'to_delete': [np.array([2]),np.array([0,2])]})

# Solution: 
df2 = df.copy()
A = list(starmap(lambda v,l: np.delete(v,l),
                             zip(df['data'],df['to_delete'])))

df2['data'] = pd.DataFrame(zip(A))
df2

prints out:

        data to_delete
0  [1, 2, 4]       [2]
1        [1]    [0, 2]
user_66303970
Related