熊猫-根据另一个列表中的索引对列表中的值求和

时间:2019-03-20 09:21:59

标签: python arrays pandas list numpy

由于我要处理大量数据,因此我试图找到最有效的方法来在短时间内解决我的问题。我的问题如下:

我有两个列表

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

但是后来我被堆积了。我正在考虑使用一些有线的循环。你有什么主意吗?

4 个答案:

答案 0 :(得分:4)

使用numpy

import numpy as np

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

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

答案 1 :(得分:3)

使用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屏蔽方法进行比较:

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

在较大的阵列上,由于bincount的工作量实际上是其工作量的2倍,因此掩蔽会变得更快。仍然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

答案 2 :(得分:0)

这很容易在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 ctor的data arg构造一个DataFrame参数:

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

然后,我们仅使用loc屏蔽df,在其中选择'b' == 'bar'的行,然后选择列'a'并调用sum()

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

答案 3 :(得分:0)

使用:

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