Как с помощью assign добавить в фрейм данных столбец, в котором будут индексы трех максимальных элементов из другого столбца той же записи?

Question

[Goblin1101]

У меня есть фрейм данных genres1, в котором у каждого элемента есть список из 20 чисел и столбец max_elements, в котором 3 максимальных числа из списка. Мне нужно сделать ещё три столбца, в которых будут индексы этих максимальных элементов. Грубо говоря нужно в списке из genres1 найти элемент, который равен max_elements[0] добавить его условно в index1 и так далее. Я понимаю что это с assign и apply делается, но раз за разом ловлю ступор мозговины и не могу додуматься, как это сделать.

genres1 = one_genre.assign(
    max_elements = lambda x: x['clusters'].apply(lambda x: heapq.nlargest(3,x)),
    cluster_index = lambda y: y['clusters'].apply(lambda x: x.index('''Затупор тут'''))
)

Answer 1

[Максим Припадчев]

Все довольно просто, на будущее если запутался, то предоставляй:
1. Входные данные (фрейм)
2 . Что хотелось бы получить.
3. Ну и попытку решения (ее ты предоставил, пытался вопросов нет).

Сделаю за тебя фрейм со случайными данными и затем решу вопрос.

df = pd.DataFrame({
    'A':[np.random.randint(1,100,20) for i in range(5)],
})

Решение

result = (
    df.assign(
        three_largest = lambda x: x['A'].apply(lambda x: pd.Series(x).nlargest(3).tolist()),
        three_largest_index = lambda x: x['A'].apply(lambda x: pd.Series(x).nlargest(3).index.tolist())
    )
)

Если захочешь разложить свои списки на отдельные колонки то используй следующий трюк.

result[['one','two','three']] = pd.DataFrame(result['three_largest_index'].tolist(),index=result.index)
result

Как видишь все в разных вариациях, как удобно.