改造一个pandas dataframe：需要更高效的解决方案

Question

I have a dataframe indexed by dates from a certain period. My columns are predictions about the value of a variable by the end of a given year. My original dataframe looks something like this:

            2016  2017  2018
2016-01-01   0.0     1   NaN
2016-07-01   1.0     1   4.1
2017-01-01   NaN     5   3.0
2017-07-01   NaN     2   2.0

where NaN means that the prediction does not exist for that given year.

Since I am working with 20+ years and most predictions are for the next 2-3 years, my real dataframe has 20+ columns mostly containing NaN values. For instance, the column for the year 2005 has predictions made in 2003-2005, but in the range 2006-2020 it's all NaN .

I would like to transform my dataframe to something like this:

            Y_0  Y_1  Y_2
2016-01-01    0    1  NaN
2016-07-01    1    1  4.1
2017-01,01    5    3  NaN
2017-07-01    2    2  NaN

where Y_j represents the prediction made for the year = index.year + j . This way, I would have a dataframe with only 4 columns (Y_0, Y_1, Y_2, Y_3).

I actually achieved this, but in what I think it is a very inefficient way:

for i in range(4):
    df[f'Y_{i}'] = numpy.nan  # create columns [Y_0, Y_1, Y_2, Y_3]

for index, row in df.iterrows():  # iterate through each row of df
    
    for year in row.dropna().index:  # iterate through each year where a prediction exists
        
        year_diff = int(year) - index.year # get the difference between the years for which the prediction was made and when it was made (possible values: 0, 1, 2 or 3)
        
        df.loc[index, f'Y_{year_diff}'] = df.loc[index, year]  # set  the values for the columns 'Y_0', 'Y_1', 'Y_2' and 'Y_3' cell by cell.

        df = df.iloc[:, -4:]  # delete all but the new columns

For a dataframe with only 1000 rows, this is taking almost 3 seconds to run. Can anyone think of a better solution?

Answer 1

You could use melt to convert it to the long format then pivot back based on the year differences.

Using your DataFrame as an example:

df = pd.DataFrame({'date':[datetime.date(2016, 1, 1), datetime.date(2016, 7, 1),
                      datetime.date(2017, 1, 1), datetime.date(2017, 7, 1)],
             2016:[0,1,np.nan,np.nan],
             2017:[1,1,5,2],
             2018:[np.nan, 4.1, 3, 2]})
df = df.melt(id_vars = 'date', value_vars = [2016, 2017, 2018], var_name='prediction_year', value_name='prediction')

Long format:

    date        prediction_year prediction
0   2016-01-01  2016    0.0
1   2016-07-01  2016    1.0
2   2017-01-01  2016    NaN
3   2017-07-01  2016    NaN
4   2016-01-01  2017    1.0
5   2016-07-01  2017    1.0
6   2017-01-01  2017    5.0
7   2017-07-01  2017    2.0
8   2016-01-01  2018    NaN
9   2016-07-01  2018    4.1
10  2017-01-01  2018    3.0
11  2017-07-01  2018    2.0

Convert back to the desired wide format:

df['year'] = pd.to_datetime(df['date']).dt.year
df['dt'] = df['prediction_year'] - df['year']
df = df.pivot(index = 'date', columns='dt', values='prediction').dropna(axis = 1, how = 'all').add_prefix('Y_')

            Y_0 Y_1 Y_2
date            
2016-01-01  0.0 1.0 NaN
2016-07-01  1.0 1.0 4.1
2017-01-01  5.0 3.0 NaN
2017-07-01  2.0 2.0 NaN

Answer 2

Let's try stack then calculate the year difference:

# in index is not already datetime
df.index = pd.to_datetime(df.index)

df = (df.stack().reset_index()
   .assign(date_diff=lambda x: x['level_1'].astype(int) - x['level_0'].dt.year)
   .pivot(index='level_0', columns='date_diff', values=0)
   .add_prefix('Y_')
)

Output:

date_diff   Y_0  Y_1  Y_2
level_0                  
2016-01-01  0.0  1.0  NaN
2016-07-01  1.0  1.0  4.1
2017-01-01  5.0  3.0  NaN
2017-07-01  2.0  2.0  NaN