Перевод кода в numpy для лучшей производительности

Question

У меня есть это:

partial = {}
for d in devs["d"]:
    for k in a1km:
        total = len(cp[(cp["r"]==d) & (cp["s"]==k)])
        partial.update({str(d)+str(k): total})

Переменные cp и devs являются pandas фреймами данных, а a1km - это словарь, содержащий сайт и все сайты в 1 км от него (предварительно рассчитано). Вывод, который мне нужен, будет для каждого d и для каждого сайта k все записи в кадре данных cp, соответствующие запросу, хранящиеся total, поэтому:

d, k, total

Я никогда не работал с numpy, и я стараюсь учиться так быстро, как только могу, но библиотека слишком велика, чтобы я мог ее обрабатывать, учитывая соображения времени в моей лаборатории. Итак, мой вопрос: как мне «перевести» приведенный ниже код в numpy для улучшения производительности?

Answer 1

Вы можете отфильтровать кадр данных и использовать pandas .DataFrame.groupby :

tmp = cp[(cp['r'].isin(devs['d'].unique()) & (cp['s'].isin(a1km))]

result_df = tmp.groupby(['r','s']).size()

Заметьте, что это также может быть довольно медленным.

Тогда чтобы сделать это в словарь:

partial = {str(k[0]) + str(k[1]): v for k,v in result_df.to_dict().items()}