Is there a better vectorized solution to writing into a dataframe using the index and columns defined as variables in a second dataframe?

Viewed 34

I have a base dataframe, which I then want to populate at certain indexes and columns as defined by a second dataframe whose rows I want to affect change in the base df. The index of this second dataframe, df_idx, are the base_df rows I'm interested in; df_idx also contains the starting and ending columns to populate, along with value to write in. base_df looks like this:

import pandas as pd
import numpy as np

months = list(range(1, 13))
li = list(map(str, months))
cols = ['ID']
cols.extend(li)

df_base = pd.DataFrame(np.random.randint(0,1000,size=(5, 13)), columns=cols)
df_base.loc[[1,2],'1':'12'] = np.nan
df_base.loc[4,'7':'12'] = np.nan

    ID      1      2      3      4      5      6      7      8      9     10     11     12
0  328   45.0  226.0  388.0  286.0  557.0  930.0  234.0  418.0  863.0  500.0  232.0  116.0
1  340    NaN    NaN    NaN    NaN    NaN    NaN    NaN    NaN    NaN    NaN    NaN    NaN
2  865    NaN    NaN    NaN    NaN    NaN    NaN    NaN    NaN    NaN    NaN    NaN    NaN
3  293  772.0  185.0    6.0  284.0  522.0  826.0  995.0  370.0   87.0  668.0  469.0   40.0
4  313  947.0  272.0  936.0  501.0  241.0  731.0    NaN    NaN    NaN    NaN    NaN    NaN

df_idx below show that for index 1, the val 210 is to be populated along the row in base_df from column '1' to column '12':

df_idx = pd.DataFrame({'start': [1, 2, 3],
                       'end': [12, 10, 11],
                       'val':np.random.randint(0,1000,size=(1, 3))[0]},
                      index=[1,2,4])

start   end val
1   1   12  210
2   2   10  663
4   3   11  922

I could loop through the rows, which I don't want to do especialy as base_df could be >250,000 rows, like so:

for idx, row in df_idx.iterrows():
    mntStrt = str(row['start'])
    mnthEnd = str(row['end'])
    df_base.loc[idx, mntStrt:mnthEnd] = row['val']

Or what I'm leaning to now, using pandas apply function:

def inputVals(x):
    idx = x.name
    mntStrt = str(x['start'])
    mnthEnd = str(x['end'])
    df_base.loc[idx, mntStrt:mnthEnd] = x['val']

df_idx.apply(lambda x: inputVals(x), axis=1)

The resulting dataframe would look like this:

    ID  1   2   3   4   5   6   7   8   9   10  11  12
0   947 537.0   827.0   477.0   39.0    586.0   370.0   576.0   556.0   119.0   158.0   990.0   958.0
1   157 129.0   129.0   129.0   129.0   129.0   129.0   129.0   129.0   129.0   129.0   129.0   129.0
2   545 NaN 849.0   849.0   849.0   849.0   849.0   849.0   849.0   849.0   849.0   NaN NaN
3   549 835.0   205.0   158.0   499.0   451.0   887.0   145.0   6.0 518.0   385.0   34.0    613.0
4   57  673.0   55.0    925.0   925.0   925.0   925.0   925.0   925.0   925.0   925.0   925.0   NaN

I feel like there's a more efficient way to tackle this problem; any insights or critiques would be most welcome. Thanks!

1 Answers

One way is to reshape df_idx to get the new value at the right spot in a dataframe indexed like df_idx and for the wanted columns 1 to 12. To do so you can use numpy and compare the columns start and end to 1 to 12. Multiply by the val columns and set index-columns like needed. so

# set sedd for reproductibility with df_idx
np.random.seed(1)

tmp = \
pd.DataFrame(
    data = ((df_idx['start'].to_numpy()[:, None] <= np.arange(1,13))
             & (df_idx['end'].to_numpy()[:, None] >= np.arange(1,13)))
            *df_idx['val'].to_numpy()[:, None], 
    index=df_idx.index, 
    columns=li
).replace(0,np.nan) # 

print(tmp)
      1      2    3    4    5    6    7    8    9   10     11    12
1  37.0   37.0   37   37   37   37   37   37   37   37   37.0  37.0
2   NaN  235.0  235  235  235  235  235  235  235  235    NaN   NaN
4   NaN    NaN  908  908  908  908  908  908  908  908  908.0   NaN

now you can use update to set the new values in df_base

df_base.update(tmp, overwrite=True) # no need of reassignment with update 
# set overwrite = False if you only change the nan values in df_base to be updated
print(df_base)
    ID      1      2      3      4      5      6      7      8      9     10  \
0   72  767.0  905.0  715.0  645.0  847.0  960.0  144.0  129.0  972.0  583.0   
1  390   37.0   37.0   37.0   37.0   37.0   37.0   37.0   37.0   37.0   37.0   
2  398    NaN  235.0  235.0  235.0  235.0  235.0  235.0  235.0  235.0  235.0   
3  319  829.0  534.0  313.0  513.0  896.0  316.0  209.0  264.0  728.0  653.0   
4  633  456.0  542.0  908.0  908.0  908.0  908.0  908.0  908.0  908.0  908.0   

      11     12  
0  749.0  508.0  
1   37.0   37.0  
2    NaN    NaN  
3  627.0  431.0  
4  908.0    NaN  
Related