时间戳

首页 > 解决方案 > 如何使用最快的方法对数据框进行分组并将列生成为数字序列？

python - 如何使用最快的方法对数据框进行分组并将列生成为数字序列？

问题描述

如何使用最快的方法对数据框进行分组并将列生成为数字序列？我的代码步骤如下：

  1. 首先生成日期数据框date
  2. 生成code数据框
  3. 生成和的Cartesian productdfdatecode
  4. 删除多余的列['a','level_1','order']
  5. 按列分组，每组内按顺序date生成一order列values

我的问题：

  1. 这些步骤感觉太麻烦了，有没有简单的方法？
  2. 第四步如何避免生成level_1andorder列
  3. 如何优化代码，现在执行需要5秒

我的代码如下：

import pandas as pd
import numpy as np


def add_order(df):
    df = df.reset_index(drop=True).reset_index()
    df = df.drop(columns='date')
    return df


def generate_data():
    np.random.seed(202107)
    date = pd.date_range(start='20150101', end='20210723', freq='D')
    date = date.to_pydatetime()
    date = np.vectorize(lambda s: s.strftime('%Y-%m-%d'))(date)
    date = pd.DataFrame(date, columns=['date'])
    date['a'] = 1

    code = pd.DataFrame(range(50), columns=['code'])
    code['a'] = 1

    df = pd.merge(date, code, how='outer')
    df['value'] = np.random.random(len(df)) * 1000

    return df


def get_result(df):
    df = df.sort_values(by='value', ascending=False)
    df = df.groupby('date').apply(add_order)
    df = df.reset_index().sort_values(by=['date', 'code']).reset_index(drop=True)
    df = df.rename(columns={'index': 'order'})
    col = ['date', 'code', 'value', 'order']
    df = df[col]
    # print(df)
    return df


def main():
    df = generate_data()
    df = get_result(df)

%timeit main()
5.25 s ± 130 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

结果是：

              date  code       value  order
0       2015-01-01     0  227.190649     39
1       2015-01-01     1  543.938036     26
2       2015-01-01     2  175.707748     43
3       2015-01-01     3  789.146427      9
4       2015-01-01     4  585.727841     24
...            ...   ...         ...    ...
119795  2021-07-23    45   92.698866     43
119796  2021-07-23    46  111.500843     40
119797  2021-07-23    47  700.675634     12
119798  2021-07-23    48  933.134534      4
119799  2021-07-23    49  108.004811     42

标签： pythonpandasdataframe

解决方案

似乎该a列是不必要的，因此生成可以变为：

def generate_data_mod():
    np.random.seed(202107)
    df = pd.MultiIndex.from_product(
        [pd.date_range(
            start='20150101', end='20210723', freq='D'
        ).strftime('%Y-%m-%d'),
         np.arange(50)],
        names=['date', 'code']
    ).to_frame(index=False)
    df['value'] = np.random.random(len(df)) * 1000
    # df['a'] = 1  # (If it is needed)
    return df

然后我们可以使用sort_valuesby value。然后用于groupby cumcount枚举组。然后sort_index恢复顺序：

def get_result_mod(df):
    # df = df.drop(columns='a')  # If df has the a column
    df = df.sort_values(by='value', ascending=False)
    df['order'] = df.groupby('date').cumcount()
    df = df.sort_index()
    return df

健全性检查：

def main():
    df = generate_data()
    df_mod = generate_data_mod()
    # True (note df_mod has no A column)
    print(df.drop(columns='a').eq(df_mod).all(None))
    # True
    print(get_result(df).eq(get_result_mod(df_mod)).all(None))

时间信息：

生成数据大致相同（merge非常有效）：

%timeit generate_data()
21 ms ± 507 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit generate_data_mod()
20.2 ms ± 428 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

get_result这样快得多：

df = generate_data()

%timeit get_result(df)
1.77 s ± 28.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit get_result_mod(df)
51 ms ± 4.17 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

推荐阅读