说我有以下数据框
import pandas as pd
df = pd.DataFrame({ 'distance':[2.0, 3.0, 1.0, 4.0],
'velocity':[10.0, 20.0, 5.0, 40.0] })
给出数据框
distance velocity
0 2.0 10.0
1 3.0 20.0
2 1.0 5.0
3 4.0 40.0
如何计算距离列的滚动总和上的速度列的平均值?使用上面的示例,在最后N行上创建一个滚动总和,以获得最小累积距离5,然后计算这些行的平均速度。
我的目标输出将是这样的:
distance velocity rv
0 2.0 10.0 NaN
1 3.0 20.0 15.0
2 1.0 5.0 11.7
3 4.0 40.0 22.5
,其中
15.0 = (10+20)/2 (2 because 3 + 2 >= 5)
11.7 = (10 + 20 + 5)/3 (3 because 1 + 3 + 2 >= 5)
22.5 = (5 + 40)/2 (2 because 4 + 1 >= 5)
更新:在Pandas中,我的代码应该找到从当前记录返回的反向累积距离总和的索引(这样它是5或更大),然后使用该索引计算移动平均线的起点。
答案 0 :(得分:5)
不是一个特别粗略的解决方案,但听起来你想要做像
这样的事情df['rv'] = np.nan
for i in range(len(df)):
j = i
s = 0
while j >= 0 and s < 5:
s += df['distance'].loc[j]
j -= 1
if s >= 5:
df['rv'].loc[i] = df['velocity'][j+1:i+1].mean()
更新:由于这个答案,OP表示他们想要一个“有效的Pandas解决方案(例如没有循环)”。如果我们认为这意味着他们想要比上述更高效的东西,那么,或许讽刺地给出评论,首先想到的优化是避免数据框架,除非需要:
l = len(df)
a = np.empty(l)
d = df['distance'].values
v = df['velocity'].values
for i in range(l):
j = i
s = 0
while j >= 0 and s < 5:
s += d[j]
j -= 1
if s >= 5:
a[i] = v[j+1:i+1].mean()
df['rv'] = a
此外,正如@JohnE所建议的那样,numba很快就会进一步优化。虽然它对上面的第一个解决方案没有太大作用,但第二个解决方案可以使用@numba.jit
开箱即用的装饰,并立即带来好处。
pd.DataFrame({'velocity': 50*np.random.random(10000), 'distance': 5*np.random.rand(10000)})
我得到以下结果:
Method Benchmark ----------------------------------------------- Original data frame based 4.65 s ± 325 ms Pure numpy array based 80.8 ms ± 9.95 ms Jitted numpy array based 766 µs ± 52 µs
即使看起来无辜的mean
也足以甩掉numba;如果我们摆脱这种情况而转而使用
@numba.jit
def numba_example():
l = len(df)
a = np.empty(l)
d = df['distance'].values
v = df['velocity'].values
for i in range(l):
j = i
s = 0
while j >= 0 and s < 5:
s += d[j]
j -= 1
if s >= 5:
for k in range(j+1, i+1):
a[i] += v[k]
a[i] /= (i-j)
df['rv'] = a
然后基准降低到158μs±8.41μs。
现在,如果您碰巧了解df['distance']
的结构,while
循环可能会进一步优化。 (例如,如果值总是远低于5,那么从尾部剪切累积总和会更快,而不是重新计算所有内容。)