python - pandas 按组索引列表分组

Question

假设我已经得到了 dataframe 的分组索引列表，并且我想使用 groupby 或其他函数获取子数据帧。我知道我可以使用 isin 像 [df[df.index.isin(group)] for group in grouplist] 多次查询数据帧，但它似乎非常慢的。我怎样才能更有效地获得这些团体？

grouplist = [[1],[29, 30, 31],[40],[46, 47, 48, 49],[58, 59],[68, 69, 70],[99, 100, 101],[103]]

Answer 1

看来你需要loc:

[df.loc[group] for group in grouplist]

编辑:

为了检查列中的值更快，使用:

[df[np.in1d(df.A, group)] for group in grouplist]

或者:

[df[df.A.isin(group)] for group in grouplist]

时间:

np.random.seed(123)
N = 100000

df = pd.DataFrame({'A': np.random.randint(150, size=N),
                   'B':np.random.rand(N)})
#print (df)


grouplist = [[1],[29, 30, 31],[40],[46, 47, 48, 49],
             [58, 59],[68, 69, 70],[99, 100, 101],[103]]

def a(df):
    df = df.set_index('A')
    return [df.loc[group] for group in grouplist]

def b(df):
    return [df[df.A.isin(group)] for group in grouplist]

def c(df):
    return [df[np.in1d(df.A, group)] for group in grouplist]


In [84]: %timeit (a(df))
10 loops, best of 3: 117 ms per loop

In [85]: %timeit (b(df))
100 loops, best of 3: 18.3 ms per loop

In [86]: %timeit (c(df))
100 loops, best of 3: 5.44 ms per loop

警告:

性能实际上取决于数据 - DataFrame 的大小和 grouplist 中组中值的数量。