Faster way to set new df value using np.array index values in dataframe

user_6051639 0
Viewed 49

I need to set the value of a new pandas df column based on the NumPy array index, also stored in the df. This works, but it is pretty slow with a large df. Any tips on how to speed things up?

a=np.random.random((5,5))
df=pd.DataFrame(np.array([[1,1],[3,3],[2,2],[3,2]]),columns=['i','j'])
df['ij']=df.apply(lambda x: (int(x['i']-1),int(x['j']-1)),axis=1)
for idx,r in df.iterrows():
    df.loc[idx,'new']=a[r['ij']]
pandas python
2 Answers

With NumPy indexing:

inds = df[["i", "j"]].to_numpy() - 1
df["new"] = a[inds[:, 0], inds[:, 1]]

where we index into a along rows with numbers in inds' first column and columns with its second column.

to get

>>> a

array([[0.27494719, 0.17706064, 0.71306907, 0.94776026, 0.04024955],
       [0.56557293, 0.63732559, 0.12254121, 0.53177861, 0.48435987],
       [0.33299644, 0.43459935, 0.57227818, 0.96142159, 0.79794503],
       [0.80112425, 0.52816002, 0.01885327, 0.39880301, 0.51974912],
       [0.60377461, 0.24419486, 0.88203753, 0.87263663, 0.49345361]])

>>> inds

array([[0, 0],
       [2, 2],
       [1, 1],
       [2, 1]])

>>> df

   i  j       new
0  1  1  0.274947
1  3  3  0.572278
2  2  2  0.637326
3  3  2  0.434599

for the ij column, you can do df["ij"] = inds.tolist().

user_93321870

Coming from the numpy side you could reshape the indices such that they match the a shape, using ravel_multi_index:

df["new"] = np.take(a, np.ravel_multi_index([df.i -1, df.j - 1], a.shape))
user_100202830
Related