pandas: вычисление года по числу месяца

Рейтинг: 1Ответов: 1Опубликовано: 16.01.2023

Прошу помочь найти лучшее решение для следующей задачи.

Дан фрейм данных, например, такой:

df_test = pd.DataFrame(np.random.randint(low = 1, high = 12, size = (10 ,3)))
df_test = df_test.rename(columns={0: 'month', 1: 'value1', 2: 'value2'})
df_test


    m   a   b
0   8   1   8
1   4   3   8
2   1   2   5
3   3   4   1
4   5   8   2
5   9   7   6
6   9   6   9
7   3   5   10
8   2   3   1
9   8   5   2

В фрейме данных отсутствует колонка с годом и присутствует колонка месяца ('m'). Известно, что начало отсчета 2021 год. Данные вносились последовательно. Также известно, что начало следующего года видно по порядку цифр в колонке месяца, когда число меньше предшествующего, то это признак следующего года.

Нужно каждому диапазону месяцев в колонке 'm' поставить год.

Я решил задачу следующим образом:

month = 0 # для сохранения месяца
years = 2021 # начало отсчета для года
df_test['year'] = 0
for count in range(0, df_test.shape[0]):
    if df_test.loc[count, 'm'] > month: # если месяц больше сохраненного значения
        month = df_test.loc[count, 'm'] # изменяем сохраненный месяц
        df_test.loc[count, 'year'] = years
    elif df_test.loc[count, 'm'] < month: # если месяц меньше сохраненного значения
        years = years + 1                 # увеличиваем года
        month = df_test.loc[count, 'm']
        df_test.loc[count, 'year'] = years
    else:
        df_test.loc[count, 'year'] = years # только присваиваем значение

df_test


    m   a   b   year
0   8   1   8   2021
1   4   3   8   2022
2   1   2   5   2023
3   3   4   1   2023
4   5   8   2   2023
5   9   7   6   2023
6   9   6   9   2023
7   3   5   10  2024
8   2   3   1   2025
9   8   5   2   2025

Помогите улучшить код. Спасибо.

Ответы

▲ 5Принят

Формируем столбец с разностью методом diff() последовательных month, далее преобразуем его в булев массив с помощью lt(0), где True - отрицательная разность (текущий месяц меньше предыдущего), затем суммируем его нарастающим итогом через cumsum(), получая последовательность от 0 до N. Ноль соответствует точке отсчета, т.е. 2021 году, 1 - 2022 и так далее. Поэтому просто прибавляем 2021 и все.

df_test['year'] = df_test.month.diff().lt(0).cumsum() + 2021
   month  value1  value2  year
0      4       5       5  2021
1      4       3      10  2021
2      1       9      11  2022
3     10      11       6  2022
4      3       9       5  2023
5      9       5       5  2023
6      6       5       7  2024
7      3       4       5  2025
8      6       7       3  2025
9      7      11       2  2025