从 Pandas 数据帧创建分组的、堆叠的 Arrays

Question

我有以下数据框：

import pandas as pd

df = pd.DataFrame({'member': ['john', 'john', 'john', 'jake', 'jake', 'jake', 'jim', 'jim', 'jim'],
                   'age': [42, 43, 44, 35, 36, 37, 57, 58, 59],
                   'inpatient_count': [0, 1, 2, 1, 0, 0, 2, 1, 5],
                   'pcp_count': [4, 0, 6, 0, 3, 3, 0, 5, 2]})

df = df.sort_values('member')

print(df)
  member  age  inpatient_count  pcp_count
3   jake   35                1          0
4   jake   36                0          3
5   jake   37                0          3
6    jim   57                2          0
7    jim   58                1          5
8    jim   59                5          2
0   john   42                0          4
1   john   43                1          0
2   john   44                2          6

我想将df转换为按member分组/嵌套的 arrays ，如下所示，但是当运行数百万个成员时，我想要更快的东西。 我希望pd.to_numpy()会有一个 grouper 参数，但我还没有弄清楚。

import numpy as np

keep = [x for x in df.columns if x != 'member']
np.array(df.groupby('member')[keep].apply(lambda x: x.values.tolist()).tolist())

array([[[35,  1,  0],
        [36,  0,  3],
        [37,  0,  3]],

       [[57,  2,  0],
        [58,  1,  5],
        [59,  5,  2]],

       [[42,  0,  4],
        [43,  1,  0],
        [44,  2,  6]]])

Answer 1

由于您之前打扰排序，您可以使用np.arraysplit来确定在哪里切割数组。 首先确保重置索引，使其成为有序的 IntIndex。

import numpy as np

df = df.sort_values('member').reset_index(drop=True)

splits = df['member'].ne(df['member'].shift()).loc[lambda x: x].index[1:]
np.array(np.array_split(df.drop(columns='member').to_numpy(), splits))
#array([[[35,  1,  0],
#        [36,  0,  3],
#        [37,  0,  3]],
# 
#       [[57,  2,  0],
#        [58,  1,  5],
#        [59,  5,  2]],
#
#       [[42,  0,  4],
#        [43,  1,  0],
#        [44,  2,  6]]])

---

对这两种方法计时，array_split（排序后）是一个不错的收获。

import perfplot
import pandas as pd
import numpy as np

def groupby_apply(df):
    keep = [x for x in df.columns if x != 'member']
    return np.array(df.groupby('member')[keep].apply(lambda x: x.values.tolist()).tolist())

def array_split(df):
    splits = df['member'].ne(df['member'].shift()).loc[lambda x: x].index[1:]
    return np.array(np.array_split(df.drop(columns='member').to_numpy(), splits))


perfplot.show(
    setup=lambda n: pd.DataFrame({'member': np.repeat(np.arange(1, n//2+1), 2),
                                  'age': np.random.randint(1, 100, n),
                                  'inpatient_count': np.random.randint(1, 10, n),
                                  'pcp_count': np.random.randint(1, 10, n)}),
    kernels=[
        lambda df: groupby_apply(df),
        lambda df: array_split(df),
    ],
    labels=['GroupBy Apply', 'np.array_split'],
    n_range=[2 ** k for k in range(2, 17)],
    equality_check=np.allclose, 
    xlabel='N Groups'
)