Question

我需要将分组的数据帧移动一个动态数字。我可以通过应用来做到这一点，但是性能不是很好。

有没有办法做到这一点的方法吗？

以下是我想做的事的一个示例：

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)

Answer 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

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

Answer 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）以外，所有计算都是在熊猫/ numpy级别上完成的，应该很快。这个想法是手动模拟apply函数中要执行的操作。

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

2 个答案: