Быстрое умножение столбца панд на годовой коэффициент

Question

У меня есть фрейм данных с индексом datetime:

df = pd.DataFrame(
    {'test':[1, 1, 1, 1, 1, 1]},
    index=[
        '2018-01-01', '2018-01-02', '2018-01-03',
        '2019-01-03', '2019-01-02', '2020-01-02'
    ]
 )
df.index=  pd.to_datetime(df.index)

У меня есть годовой параметр:

yearly_parameter = [1, 2, 3]

Я хотел бы эффективно умножить (в векторизованном виде?) Столбец'test' по соответствующему годовому параметру, содержащемуся в списке year_parameter (первое значение для 2018 года, второе для 2019 года и третье для 2020 года).Как я могу сделать это эффективно?Является ли список хорошим способом хранения этих годовых параметров для расчета?

Я ожидаю следующий результат в столбце, скажем «ответ»:

df['answer'] = [1, 1, 1, 2, 2, 3]

print(df)

              test  answer
2018-01-01     1       1
2018-01-02     1       1
2018-01-03     1       1
2019-01-03     1       2
2019-01-02     1       2
2020-01-02     1       3

Большое спасибо за вашу помощь,

Пьер

Answer 1

pd.factorize

Используйте factorize, чтобы установить порядок лет, который должен соответствовать элементам yearly_parameter.Тогда мы можем эффективно умножить срезание массива.

Это предполагает, что длина yearly_parameter будет, по крайней мере, равна числу уникальных лет в df.index

f, y = pd.factorize(df.index.year)

yearly_parameter = np.array([1, 2, 3])

df.assign(answer=df.test.values * yearly_parameter[f])

            test  answer
2018-01-01     1       1
2018-01-02     1       1
2018-01-03     1       1
2019-01-03     1       2
2019-01-02     1       2
2020-01-02     1       3

---

np.unique

Помните, что это предполагает, что yearly_parameter выравнивает свой первый элемент с первым наблюдаемым годом.Если вы хотите, чтобы первый элемент соответствовал наблюдаемому минимальному году, вам следует использовать pd.factorize(df.index.year, sort=True).Или еще лучше, если вы собираетесь сортировать, тогда используйте эквивалентный расчет в Numpy

y, f = np.unique(df.index.year, return_inverse=True)

yearly_parameter = np.array([1, 2, 3])

df.assign(answer=df.test.values * yearly_parameter[f])

            test  answer
2018-01-01     1       1
2018-01-02     1       1
2018-01-03     1       1
2019-01-03     1       2
2019-01-02     1       2
2020-01-02     1       3