Pandas - 根據另一個列表中的索引對列表中的值求和

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'] 

但后來我堆積了。 我正在考慮使用一些有線for循環。 你有什么主意嗎？

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

比較例如numpy masking方法：

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

在較大的陣列上，掩蔽變得稍快，這是預期的，因為bincount基本上是工作的2 bincount 。 仍然bincount時間不到2倍，所以如果你碰巧需要兩個總和（'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

這在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

所以在這里我們采用你的列表並為DataFrame ctor的data arg構建一個dict ：

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

然后我們使用loc掩蓋df，我們選擇'b' == 'bar'並選擇列'a'並調用sum() ：

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

Answer 4

采用：

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']