Панды - сумма значений в списке в соответствии с индексом из другого списка

Question

Я пытаюсь найти самый питонический способ решения моей проблемы в кратчайшие сроки, так как я имею дело с большим количеством данных.Моя проблема заключается в следующем:

У меня есть два списка

a = [12,34,674,2,0,5,6,8]
b = ['foo','bar','bar','foo','foo','bar','foo','foo']

Я хочу сказать python: если 'bar' находится в b, взять все индексы и суммировать все значения в списке aс этими индексами.

Это то, что я сделал до сих пор:

idx = [i for i, j in enumerate(a) if j == 'bar'] 

, но потом меня сложили.Я рассматриваю возможность использования некоторых проводных петель.У вас есть идеи?

Answer 1

С numpy:

import numpy as np

a = np.array(a)
b = np.array(b)

a[b == 'bar'].sum()

Answer 2

Использование np.bincount.Вычисляет обе суммы ('foo' и 'bar').

sum_foo, sum_bar = np.bincount(np.char.equal(b, 'bar'), a)
sum_foo
# 28.0
sum_bar
# 713.0

Примечание np.char.equal работает как со списками, так и с массивами.Если b является массивом, тогда вместо него можно использовать b == 'bar' и он немного быстрее.

Время:

Несмотря на то, что это вычисляет обе суммы, на самом деле это довольно быстро:

timeit(lambda: np.bincount(b == 'bar', a))
# 2.406161994993454

Сравните, например, с методом маскировки пустышек:

timeit(lambda: a[b == 'bar'].sum())
# 5.642918559984537

На больших массивах маскирование становится немного быстрее, что ожидается, поскольку bincount выполняет в два раза большую работу.Тем не менее, bincount занимает меньше, чем двукратное время, поэтому, если вам понадобятся обе суммы ('foo' и 'bar'), bincount все еще быстрее.

aa = np.repeat(a, 1000)
bb = np.repeat(b, 1000)
timeit(lambda: aa[bb == 'bar'].sum(), number=1000)
# 0.07860603698645718
timeit(lambda:np.bincount(bb == 'bar', aa), number=1000)
# 0.11229897901648656

Answer 3

Использование:

l = [x for x,y in zip(a,b) if y == 'bar']

Если вы хотите индексы:

l = [i for (i,x),y in zip(enumerate(a),b) if y == 'bar']

Answer 4

Это просто сделать в pandas:

In[5]:
import pandas as pd
a = [12,34,674,2,0,5,6,8]
b = ['foo','bar','bar','foo','foo','bar','foo','foo']
df = pd.DataFrame({'a':a, 'b':b})
df

Out[5]: 
     a    b
0   12  foo
1   34  bar
2  674  bar
3    2  foo
4    0  foo
5    5  bar
6    6  foo
7    8  foo

In [8]: df.loc[df['b']=='bar','a'].sum()
Out[8]: 713

Итак, мы возьмем ваши списки и создадим dict вместо аргумента data для DataFrame ctor:

df = pd.DataFrame({'a':a, 'b':b})

Затем мы просто маскируем df, используя loc, где мы выбираем строки, где 'b' == 'bar', выбираем столбец 'a' и вызываем sum():

df.loc[df['b']=='bar','a'].sum()