如何基于熊猫数据框中的其他行创建新列?

时间:2020-01-03 13:02:18

标签: python pandas

我有一个包含20万行的数据框,并且在某些情况下,我尝试根据其他行添加列。我试图实现这一目标,但要花很多时间(2小时)。

这是我的代码:

for index in dataset.index:
    A_id = dataset.loc[index, 'A_id']
    B_id = dataset.loc[index, 'B_id']
    C_date = dataset.loc[index, 'C_date']
    subset = dataset[
        (dataset['A_id'] == A_id) & (dataset['B_id'] == B_id) & (
                dataset['C_date'] < C_date)]
    dataset.at[index, 'D_mean'] = subset['D'].mean()
    dataset.at[index, 'E_mean'] = subset['E'].mean()

我的数据框看起来像这样:

A = [1, 2, 1, 2, 1, 2]
B = [10, 20, 10, 20, 10, 20]
C = ["22-02-2019", "28-02-19", "07-03-2019", "14-03-2019", "21-12-2019", "11-10-2019"]
D = [10, 12, 21, 81, 20, 1]
E = [7, 10, 14, 31, 61, 9]

dataset = pd.DataFrame({
    'A_id': A,
    'B_id': B,
    'C_date': C,
    'D': D,
    'E': E,
})

dataset.C_date = pd.to_datetime(dataset.C_date)
dataset
Out[27]: 
   A_id  B_id     C_date   D   E
0     1    10 2019-02-22  10   7
1     2    20 2019-02-28  12  10
2     1    10 2019-07-03  21  14
3     2    20 2019-03-14  81  31
4     1    10 2019-12-21  20  61
5     2    20 2019-11-10   1   9

我想以比我的解决方案更好的有效方式获得此结果:

   A_id  B_id     C_date   D   E  D_mean  E_mean
0     1    10 2019-02-22  10   7     NaN     NaN
1     2    20 2019-02-28  12  10     NaN     NaN
2     1    10 2019-07-03  21  14    10.0     7.0
3     2    20 2019-03-14  81  31    12.0    10.0
4     1    10 2019-12-21  20  61    15.5    10.5
5     2    20 2019-11-10   1   9    46.5    20.5

你有个主意吗?

3 个答案:

答案 0 :(得分:3)

我们可以使用功能组合来实现此目的,最著名的是pd.DataFrame.rolling来计算移动平均值。

def custom_agg(group):
    cols = ['D', 'E']
    for col in cols:
        name = '{}_mean'.format(col)
        group[name] = group[col].shift() \
                                .rolling(len(group[col]), min_periods=2) \
                                .mean() \
                                .fillna(group[col].iloc[0])
        group[name].iloc[0] = pd.np.nan
    return group

dataset.groupby(['A_id', 'B_id'], as_index=False).apply(custom_agg)

   A_id  B_id     C_date   D   E  D_mean  E_mean
0     1    10 2019-02-22  10   7     NaN     NaN
1     2    20 2019-02-28  12  10     NaN     NaN
2     1    10 2019-07-03  21  14    10.0     7.0
3     2    20 2019-03-14  81  31    12.0    10.0
4     1    10 2019-12-21  20  61    15.5    10.5
5     2    20 2019-11-10   1   9    46.5    20.5

可能会有更优雅的方法来执行此操作,但是使用此方法可以提高性能。只需确保C_date列已提前移动即可,因为它是移动平均数。

答案 1 :(得分:0)

我怀疑您在循环中创建子集的成本很高,而我的测试表明,您的算法以每分钟约11,000个索引的速度运行。我想出了一种替代算法,可以对数据进行预排序,以使计算子集变得微不足道,并且在20万行随机数据集上运行需要不到5分钟的时间。

dataset.sort_values(by=['A_id', 'B_id', 'C_date'], inplace=True)
dataset.reset_index(drop=True, inplace=True)

last_A = None
last_B = None
first_index = -1
for index in dataset.index:
    A_id = dataset.loc[index, 'A_id']
    B_id = dataset.loc[index, 'B_id']
    C_date = dataset.loc[index, 'C_date']

    if (last_A != A_id) | (last_B != B_id):
        first_index = index
        last_A = A_id
        last_B = B_id

    subset = dataset[first_index:index]
    dataset.at[index, 'D_mean'] = subset['D'].mean()
    dataset.at[index, 'E_mean'] = subset['E'].mean()

答案 2 :(得分:0)

这是使用.apply的一种方法:

dataset[['D_mean', 'E_mean']] = (dataset
                                .apply(lambda df: dataset[(dataset['A_id'] == df['A_id']) & 
                                                          (dataset['B_id'] == df['B_id']) & 
                                                          (dataset['C_date'] < df['C_date'])
                                                          ][['D','E']].mean(axis=0), axis=1)

   A_id  B_id     C_date   D   E  D_mean  E_mean
0     1    10 2019-02-22  10   7     NaN     NaN
1     2    20 2019-02-28  12  10     NaN     NaN
2     1    10 2019-07-03  21  14    10.0     7.0
3     2    20 2019-03-14  81  31    12.0    10.0
4     1    10 2019-12-21  20  61    15.5    10.5
5     2    20 2019-11-10   1   9    46.5    20.5