Более быстрый способ построения мультииндекс даты в Pandas - PullRequest
Новая
       2

Более быстрый способ построения мультииндекс даты в Pandas

0 голосов
/ 26 февраля 2020

У меня есть Pandas фрейм данных, df. Вот первые пять строк: 

    Id  StartDate   EndDate
0   0   2015-08-11  2018-07-13
1   1   2014-02-15  2016-01-25
2   2   2014-12-20  NaT
3   3   2015-01-09  2015-01-14
4   4   2014-07-20  NaT

Я хочу создать новый фрейм данных, df2. У df2 должна быть строка для каждого месяца от StartDate до EndDate включительно, для каждого Id в df1. Например, поскольку первая строка df1 имеет StartDate в августе 2015 года и EndDate в июле 2018 года, в df2 должны быть строки, соответствующие августу 2015 года, сентябрю 2015 года, ..., июлю 2018. Если значение Id в df1 не имеет EndDate, мы примем его за июнь 2019 года.

Я бы хотел df2 использовать мультииндекс с первым уровнем, соответствующим Id в df1 второй уровень - год, а третий уровень - месяц. Например, если все вышеперечисленные пять строк были df1, то df2 должно выглядеть следующим образом: 

Id  Year    Month
0   2015    8
            9
            10
            11
            12
    2016    1
            2
            3
            4
            5
            6
            7
            8
            9
            10
            11
            12
    2017    1
            2
            3
            4
            5
            6
            7
            8
            9
            10
            11
            12
    2018    1
... ... ...
4   2017    1
            2
            3
            4
            5
            6
            7
            8
            9
            10
            11
            12
    2018    1
            2
            3
            4
            5
            6
            7
            8
            9
            10
            11
            12
    2019    1
            2
            3
            4
            5
            6

Следующий код справляется с задачей, но на моем приличном ноутбуке занимает около 20 секунд за 10 тыс. Id s. Могу ли я быть более эффективным как-то? 

import numpy as np

def build_multiindex_for_id_(id_, enroll_month, enroll_year, cancel_month, cancel_year):
    # Given id_ and start/end dates,
    # returns 2d array to be converted to multiindex.
    # Each row of returned array represents a month/year
    # between enroll date and cancel date inclusive.
    year = enroll_year
    month = enroll_month
    multiindex_array = [[],[],[]]
    while (month != cancel_month) or (year != cancel_year):
        multiindex_array[0].append(id_)
        multiindex_array[1].append(year)
        multiindex_array[2].append(month)
        month += 1
        if month == 13:
            month = 1
            year += 1
    multiindex_array[0].append(id_)
    multiindex_array[1].append(year)
    multiindex_array[2].append(month)    
    return np.array(multiindex_array)


# Begin by constructing array for first id.
array_for_multiindex = build_multiindex_for_id_(0,8,2015,7,2018)

# Append the rest of the multiindices for the remaining ids.
for _, row in df.loc[1:].fillna(pd.to_datetime('2019-06-30')).iterrows():
    current_id_array = build_multiindex_for_id_(
        row['Id'],
        row['StartDate'].month,
        row['StartDate'].year,
        row['EndDate'].month,
        row['EndDate'].year)
    array_for_multiindex = np.append(array_for_multiindex, current_id_array, axis=1)

df2_index = pd.MultiIndex.from_arrays(array_for_multiindex).rename(['Id','Year','Month'])

pd.DataFrame(index=df2_index)

1 Ответ

1 голос
/ 26 февраля 2020

Вот мой подход после нескольких проб и ошибок: 

(df.melt(id_vars='Id')
   .fillna(pd.to_datetime('June 2019'))
   .set_index('value')
   .groupby('Id').apply(lambda x: x.asfreq('M').ffill())
   .reset_index('value')
   .assign(year=lambda x: x['value'].dt.year,
           month=lambda x: x['value'].dt.month)
   .set_index(['year','month'], append=True)
)

Вывод: 

                   value  Id variable
Id year month                        
0  2015 8     2015-08-31 NaN      NaN
        9     2015-09-30 NaN      NaN
        10    2015-10-31 NaN      NaN
        11    2015-11-30 NaN      NaN
        12    2015-12-31 NaN      NaN
   2016 1     2016-01-31 NaN      NaN
        2     2016-02-29 NaN      NaN
        3     2016-03-31 NaN      NaN
        4     2016-04-30 NaN      NaN
        5     2016-05-31 NaN      NaN
        6     2016-06-30 NaN      NaN
        7     2016-07-31 NaN      NaN
        8     2016-08-31 NaN      NaN
        9     2016-09-30 NaN      NaN
        10    2016-10-31 NaN      NaN
...