Как эффективно выбрать первые N столбцов, группируя их для каждой строки DataFrame pandas?

import numpy as np
import pandas as pd

rng = np.random.default_rng()
dates = pd.date_range('2024-08-01', '2024-08-07')
contestants = ['Alligator', 'Beryl', 'Chupacabra', 'Dandelion', 'Eggplant', 'Feldspar']
coolness_score = pd.DataFrame(rng.random((len(dates), len(contestants))), index=dates, columns=contestants)

## solution:
# melt dataframe(from wide to long data format)
df = coolness_score.reset_index().rename(columns = {"index": "date"}).melt(id_vars = "date", var_name = "contestant", value_name = "coolness_score")

# map contestant to category and insert category column after date column
category_mapping = {
    'Alligator': 'Animal',
    'Beryl': 'Mineral',
    'Chupacabra': 'Animal',
    'Dandelion': 'Vegetable',
    'Eggplant': 'Vegetable',
    'Feldspar': 'Mineral'
}
df.insert(1, "category", df["contestant"].map(category_mapping))

# sort by date and category for better readability
df = df.sort_values(by=["date", "category"], ignore_index=True)

# the id of the best single score from each category for each day
day_category_top = df.groupby(["date", "category"])[["coolness_score"]].idxmax().rename(columns = {"coolness_score": "best_score_index"})

# 1. The best single score from each category
df.loc[day_category_top["best_score_index"]]

# 2. The best N scores from each category, where N is consistent across all categories
N = 3
df.groupby(["date", "category"]).apply(lambda g: list(g["coolness_score"].nlargest(N)))

это решение предполагает, что в определенный день каждый участник появляется один раз (т. е. один балл в день на каждого участника), о чем и предполагают ваши данные.

для третьего сценария было неясно, имеете ли вы в виду применить ограничение к общему баллу категории или нет. если так:

# 3. The best scores from each category with limits defined by category_limits
category_limits = {
    'Animal': 1,
    'Vegetable': 2,
    'Mineral': 1
}
# sum of the scores for each category, but capped at limit
sum_score = df.groupby(["date", "category"])[["coolness_score"]].sum().reset_index()
sum_score["actual_score"] = sum_score.apply(lambda x: min(x["coolness_score"], category_limits[x["category"]]), axis=1)
sum_score

и вы можете применить аналогичный логин к 1 и 2 к actual_score.

Вот возможное решение данной проблемы. Сценарий 1: Лучший одиночный результат в каждой категории В первом сценарии мы хотим получить лучший отдельный балл в каждой категории за каждую дату.

def get_top_score_per_category(row):
    grouped = row.groupby(category_mapping, axis=0)
    return grouped.apply(lambda x: x.idxmax())

top_score_indices = coolness_score.apply(get_top_score_per_category, axis=1)
top_scores = coolness_score.lookup(coolness_score.index, top_score_indices)

Это позволит определить лучшего участника в каждой категории на каждый день.

Сценарий 2: N лучших результатов в каждой категории Если нам нужны N лучших результатов в каждой категории:

N = 2  # Example for the top 2 scores

def get_top_n_scores_per_category(row):
    grouped = row.groupby(category_mapping, axis=0)
    return grouped.apply(lambda x: x.nlargest(N).index)

top_n_scores_indices = coolness_score.apply(get_top_n_scores_per_category, axis=1)

Это возвращает N лучших участников в каждой категории на каждый день.

Сценарий 3: Лучшие результаты в каждой категории с ограничениями Теперь давайте реализуем самый сложный сценарий, в котором у вас есть ограничения, определенные с помощью Category_limits.

def get_top_scores_with_limits(row):
    grouped = row.groupby(category_mapping, axis=0)
    top_scores = []
    for category, group in grouped:
        limit = category_limits[category]
        top_scores.extend(group.nlargest(limit).index)
    return top_scores

top_scores_with_limits_indices = coolness_score.apply(get_top_scores_with_limits, axis=1)

Это даст вам индексы участников, соответствующих критериям, основанным на ограничениях категорий.

Обнуление очков проигравших Чтобы обнулить баллы проигравших (участников, не соответствующих критериям):

def zero_out_losers(row):
    top_scores = get_top_scores_with_limits(row)
    return row.where(row.index.isin(top_scores), 0)

final_scores = coolness_score.apply(zero_out_losers, axis=1)

Это возвращает DataFrame, в котором остаются только самые высокие оценки, а все остальные оценки устанавливаются на ноль.

Дайте мне знать, если это сработает для вас.

Вы можете использовать сопоставление и groupby.rank , затем маску с , где:

# names to categories
cat = coolness_score.columns.map(category_mapping)
# categories to limits
limit = cat.map(category_limits)

# rank per category
rank = coolness_score.T.groupby(cat).rank(method='dense', ascending=False).T
# identify top N per category per row
mask = rank.le(limit)

# mask losers
out = coolness_score.where(mask, 0)

Выход:

            Alligator     Beryl  Chupacabra  Dandelion  Eggplant  Feldspar
2024-08-01   0.000000  0.878593    0.957980   0.887114  0.266656  0.000000
2024-08-02   0.000000  0.660319    0.737451   0.921197  0.446438  0.000000
2024-08-03   0.000000  0.000000    0.765396   0.334504  0.250021  0.736392
2024-08-04   0.000000  0.000000    0.990308   0.357501  0.124491  0.941783
2024-08-05   0.327078  0.000000    0.000000   0.309475  0.538202  0.952041
2024-08-06   0.533576  0.000000    0.000000   0.935781  0.587427  0.690166
2024-08-07   0.767592  0.000000    0.000000   0.222281  0.879662  0.821808

Промежуточные продукты:

# cat
Index(['Animal', 'Mineral', 'Animal', 'Vegetable', 'Vegetable', 'Mineral'], dtype='object')

# limit
Index([1, 1, 1, 2, 2, 1], dtype='int64')

# rank
            Alligator  Beryl  Chupacabra  Dandelion  Eggplant  Feldspar
2024-08-01        2.0    1.0         1.0        1.0       2.0       2.0
2024-08-02        2.0    1.0         1.0        1.0       2.0       2.0
2024-08-03        2.0    2.0         1.0        1.0       2.0       1.0
2024-08-04        2.0    2.0         1.0        1.0       2.0       1.0
2024-08-05        1.0    2.0         2.0        2.0       1.0       1.0
2024-08-06        1.0    2.0         2.0        1.0       2.0       1.0
2024-08-07        1.0    2.0         2.0        2.0       1.0       1.0

# mask
            Alligator  Beryl  Chupacabra  Dandelion  Eggplant  Feldspar
2024-08-01      False   True        True       True      True     False
2024-08-02      False   True        True       True      True     False
2024-08-03      False  False        True       True      True      True
2024-08-04      False  False        True       True      True      True
2024-08-05       True  False       False       True      True      True
2024-08-06       True  False       False       True      True      True
2024-08-07       True  False       False       True      True      True