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

Question

У меня есть датафрейм с именем (person_name), цветом (shirt_color) в качестве столбцов Каждый человек в определенный день носит рубашку определенного цвета (количество дней может быть произвольным)

например, ввод:

name    color
----------------
John    White
John    White
John    Blue
John    Blue
John    White
Tom     Blue
Tom     Blue
Tom     Green
Tom     Black
Jerry   Black
Jerry   Blue
Jerry   Black

Мне нужно найти лучшую цветную рубашку, которую носит каждый человек, в соответствии с best_color_order например, результат:

name    color
-------------
Jerry   Blue
John    White
Tom     Green

Я делаю следующее:

best_color_order = ['White', 'Green', 'Blue', 'Black']

best_color_list = [[name, list(group['color']).sort(key=best_color_order.index)[0]]
                    for name, group in df.groupby(by='name', sort=False, as_index=False)]

best_color_df = pd.DataFrame(best_color_list, columns=['name', 'color'])

Есть ли более быстрый способ сделать это, если у меня есть миллионы записей?

Answer 1

Преобразуйте столбец color в упорядоченную категориальную , чтобы она сортировалась в желаемом порядке, затем отсортируйте значения по color и, наконец, groupby и примите значение first для каждое имя:

best_color_order = ['White', 'Green', 'Blue', 'Black']

df['color'] = pd.Categorical(df['color'], categories = best_color_order, ordered=True)

df.sort_values('color').groupby('name').first()

       color
name        
Jerry   Blue
John   White
Tom    Green

[EDIT] : более быстрым способом может быть то же самое, но вместо группового режима просто отбросьте дубликат name s и оставьте первый (по умолчанию для функции drop_duplicates):

best_color_order = ['White', 'Green', 'Blue', 'Black']

df['color'] = pd.Categorical(df['color'], categories = best_color_order, ordered=True)

df.sort_values('color').drop_duplicates('name')

     name  color
0    John  White
7     Tom  Green
10  Jerry   Blue