pandas 根据列值指定数量向上求部分数据和

看过来

《pandas 教程》持续更新中，提供建议、纠错、催更等加作者微信: gairuo123（备注：pandas教程）和关注公众号「盖若」ID: gairuo。跟作者学习，请进入 Python学习课程。欢迎关注作者出版的书籍：《深入浅出Pandas》和《Python之光》。

pandas 中支持对数据的切片和索引，这个机制让我们选择子Series非常方便。在本例中，我们将实现一个由相应数字控制、向上取部分Series的需求。

数据与需求

有以下构造的数据：

import pandas as pd
import random

random.seed(666)

df = pd.DataFrame({'a': random.choices(range(10), k=6),
                   'b': random.choices(range(1, 4), k=6)}
                 )
df
'''
   a  b
0  4  3
1  9  1
2  4  3
3  5  2
4  8  2
5  5  3
'''

a列为基础数据，b列为一个标记，代码需要相加对应a列的及以前的几个数字，将新计算的数据增加到后边一列。

比如：索引3位的b值为2，需要将包含a当前行往上两个值相加，5加4为9。

思路

首先要注意的是，需求中的算法不是累加，不能直接用 cumsum() 计算。

最简单的办法是写一个函数传入b所在的行，来对DataFrame进行计算，然后用 apply() 或者 map() 来应用这个函数。

关于计算逻辑也非常清楚，传入整行，以该行所在的索引截取b值所在行以上的所有DataFrame部分，然后向上取b数量个部分，最后相加起来。

代码

我们先以索引3位的为例我们写单个处理逻辑。先查询出索引3的行数据：

df.loc[3]
'''
a    5
b    2
Name: 3, dtype: int64
'''

可以看到得以整行的Series的name是3，取3以上的部分的a列：

df.a.loc[:3]
'''
0    4
1    9
2    4
3    5
Name: a, dtype: int64
'''

b的值是2，取从下往上的两个值：

df.a.loc[:3][-2:]
'''
2    4
3    5
Name: a, dtype: int64
'''

也可以先对这部分Series反转再取前两个：

df.a[3::-1].loc[:2]
'''
3    5
2    4
Name: a, dtype: int64
'''

最后对它们求和：

df.a[3::-1].loc[:2].sum()
# 9

将以上逻辑写一个 lambda 匿名函数，最终的代码是：

df.assign(
    cumsum_b = 
        df.apply(lambda x: df.a[x.name::-1].iloc[:int(x.b)].sum(),
         axis=1)
)
'''
   a  b  cumsum_b
0  4  3         4
1  9  1         9
2  4  3        17
3  5  2         9
4  8  2        13
5  5  3        18
'''

这样，我们就完成了这个需求。

（完）

pandas 根据列值指定数量向上求部分数据和

数据与需求

思路

代码

相关内容