在数据帧上使用 groupby 进行动态移位答案

【问题标题】：dynamic shift with groupby on dataframe在数据帧上使用 groupby 进行动态移位
【发布时间】：2019-06-30 20:50:23
【问题描述】：

我需要将分组数据框移动一个动态数字。我可以用apply来做，但是性能不是很好。

有什么方法可以在没有应用的情况下做到这一点？

这是我想做的一个示例：

df = pd.DataFrame({
    'GROUP': ['A', 'A', 'A', 'A', 'A', 'A', 'B','B','B','B','B','B'], 
    'VALUE': [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2],
    'SHIFT': [ 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3]
    })  

df['SUM'] = df.groupby('GROUP').VALUE.cumsum()

# THIS DOESN'T WORK:
df['VALUE'] = df.groupby('GROUP').SUM.shift(df.SHIFT)

我通过以下方式应用：

df = pd.DataFrame({
    'GROUP': ['A', 'A', 'A', 'A', 'A', 'A', 'B','B','B','B','B','B'], 
    'VALUE': [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2],
    'SHIFT': [ 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3]
    })  

def func(group):
    s = group.SHIFT.iloc[0]

    group['SUM'] = group.SUM.shift(s)

    return group

df['SUM'] = df.groupby('GROUP').VALUE.cumsum()

df = df.groupby('GROUP').apply(func)

【问题讨论】：

我只需要一种更快的方法，因为即使使用 Dask，在数百万行上运行的应用程序也表现不佳。

标签： python python-3.x pandas

【解决方案1】：

如果数据框按组排序（如您的示例），则这是一个纯 numpy 版本：

# these rows are not null after shifting
notnull = np.where(df.groupby('GROUP').cumcount() >= df['SHIFT'])[0]
# source rows for rows above
source = notnull - df['SHIFT'].values[notnull]

shifted = np.empty(df.shape[0])
shifted[:] = np.nan
shifted[notnull] = df.groupby('GROUP')['VALUE'].cumsum().values[source]
df['SUM'] = shifted

它首先获取要更新的行的索引。可以减去移位以产生源行。

【讨论】：

我还有一些未在我的问题中定义的要求，并且由于我能够使用 join 以另一种方式做到这一点，因此我会将您的回答正确，因为它很好地回答了我的问题.谢谢！

【解决方案2】：

如果组是连续的，则避免apply 的解决方案可能如下：

import numpy as np
import pandas as pd

df = pd.DataFrame({
    'GROUP': ['A', 'A', 'A', 'A', 'A', 'A', 'B','B','B','B','B','B'],
    'VALUE': [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2],
    'SHIFT': [ 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3]
    })


# compute values required for the slices
_, start = np.unique(df.GROUP.values, return_index=True)
gp = df.groupby('GROUP')
shifts = gp.SHIFT.first()
sizes = gp.size().values
end = (sizes - shifts.values) + start

# compute slices
source = [i for s, f in zip(start, end) for i in range(s, f)]
target = [i for j, s, f in zip(start, shifts, sizes) for i in range(j + s, j + f)]

# compute cumulative sum and arrays of nan
s = gp.VALUE.cumsum().values
r = np.empty_like(s, dtype=np.float32)
r[:] = np.nan

# set the on the array of nan
np.put(r, target, s[source])

# set the sum column
df['SUM'] = r

print(df)

输出

   GROUP  SHIFT  VALUE   SUM
0      A      2      1   NaN
1      A      2      2   NaN
2      A      2      3   1.0
3      A      2      4   3.0
4      A      2      5   6.0
5      A      2      6  10.0
6      B      3      7   NaN
7      B      3      8   NaN
8      B      3      9   NaN
9      B      3      0   7.0
10     B      3      1  15.0
11     B      3      2  24.0

除了构建切片（source 和 target）之外，所有计算都在应该很快的 pandas/numpy 级别完成。这个想法是手动模拟应用函数中的操作。

【讨论】：