Pandas - flatten the diagonal values of dataframe returned from an apply loc operation

Question

I have two dataframes -- for each row in df I want to look up the matching epoch_minute - lag and grab the corresponding average_hc value.

>>> df.head()
   epoch_minute  headcount
0      25640940          8
1      25640939          7
2      25640938          6
3      25640937         10
4      25640936         11
>>> avgs.head()
    epoch_minute  average_hc
0      25640940         7.8
1      25640939         8.5
2      25640938         9.2
3      25640937        10.0
4      25640936        10.1

This for-loop does the trick:

lag = []
for i,r in df.iterrows():
    l = avgs.loc[avgs['epoch_minute'] == (r['epoch_minute'] - day), 'average_hc'] 
    lag.append(l) if not l.empty else np.nan       
avgs['lag'] = pd.Series(lag).astype('float')

It produces this correct Series lag:

  epoch_minute  average_hc   lag  headcount
0      25640940         7.8  30.2        8.0
1      25640939         8.5  28.1        7.0
2      25640938         9.2  26.2        6.0
3      25640937        10.0  24.7       10.0
4      25640936        10.1  23.1       11.0

When I try the same logic as an apply operation:

lag1 = df.apply(lambda r: avgs.loc[avgs['epoch_minute'] == (r['epoch_minute'] - day),
         'average_hc'], axis=1)
print(f"lag1.shape: {lag1.shape}\nlag1.head():\n{lag1.head()}")

It produces this correctly valued but incorrectly shaped dataframe instead of the Series I expected:

        1285  1286       1287       1288  1289  1290  1291  1292  1293  1294  ...   2655  2656  2657  2658  2659  2660  2661  2662  2663  2664
0  30.200001   NaN        NaN        NaN   NaN   NaN   NaN   NaN   NaN   NaN  ...    NaN   NaN   NaN   NaN   NaN   NaN   NaN   NaN   NaN   NaN
1        NaN  28.1        NaN        NaN   NaN   NaN   NaN   NaN   NaN   NaN  ...    NaN   NaN   NaN   NaN   NaN   NaN   NaN   NaN   NaN   NaN
2        NaN   NaN  26.299999        NaN   NaN   NaN   NaN   NaN   NaN   NaN  ...    NaN   NaN   NaN   NaN   NaN   NaN   NaN   NaN   NaN   NaN
3        NaN   NaN        NaN  24.799999   NaN   NaN   NaN   NaN   NaN   NaN  ...    NaN   NaN   NaN   NaN   NaN   NaN   NaN   NaN   NaN   NaN
4        NaN   NaN        NaN        NaN  23.1   NaN   NaN   NaN   NaN   NaN  ...    NaN   NaN   NaN   NaN   NaN   NaN   NaN   NaN   NaN   NaN

How can I remove the nans or flatten the dataframe that is returned by the apply so it is a Series of the diagonal values?

If you want to use it as a new column in the dataframe it would be easier to use `merge`. — Shaido, Oct 02 '18 at 05:56

jezrael · Accepted Answer · 2018-10-02T06:10:37.410

Use next with iter for return default value NaN if no value matched and is returneded empty Series:

ag1 = df.apply(lambda r: next(iter(avgs.loc[avgs['epoch_minute'] == (r['epoch_minute'] - day),
         'average_hc']), np.nan), axis=1)

Verify:

day = 1
lag = []
for i,r in df.iterrows():
    l = avgs.loc[avgs['epoch_minute'] == (r['epoch_minute'] - day), 'average_hc'] 
    lag.append(l) if not l.empty else np.nan       
avgs['lag'] = pd.Series(lag).astype('float')

avgs['lag1'] = df.apply(lambda r: next(iter(avgs.loc[avgs['epoch_minute'] == 
                                          (r['epoch_minute'] - day), 'average_hc']), np.nan),
                                       axis=1)
print (avgs)
   epoch_minute  average_hc   lag  lag1
0      25640940         7.8   8.5   8.5
1      25640939         8.5   9.2   9.2
2      25640938         9.2  10.0  10.0
3      25640937        10.0  10.1  10.1
4      25640936        10.1   NaN   NaN

Another solution with merge and left join, last shift lag column by shift:

df1 = df[['epoch_minute']]
avgs1 = avgs.rename(columns={'average_hc':'lag1'})
df2 = avgs1.merge(df1, how='left', on='epoch_minute')
df2['lag1'] = df2['lag1'].shift(-day)

Pandas - flatten the diagonal values of dataframe returned from an apply loc operation

1 Answers1