python - 基于字典的 Pandas 示例

Question

我正在尝试根据字典和特定列对 pandas DataFrame 进行采样。因此，对于 y 列的每个值，我确切地知道我想要选择多少个观察值。

我可以通过 groupby apply 组合来做到这一点:

import pandas as pd

df = pd.DataFrame({'y': [2,2,0,0,0,1,1,1,1,1], 'x': 1, 'z': 2})

    y   x   z
0   2   1   2
1   2   1   2
2   0   1   2
3   0   1   2
4   0   1   2
5   1   1   2

sizes = {0: 2, 1: 1, 2:1}

df.groupby('y').apply(lambda x: x.sample(sizes[x['y'].values[0]]))

yyxz

0 2 0 1 2 4 0 1 21 5 1 1 22 0 2 1 2

但是，如果我使用 unique 而不是 values (这应该是等效的)，我会收到一个奇怪的 KeyError: 'y' 错误数据框:

df.groupby('y').apply(lambda x: x.sample(sizes[x.y.unique()[0]]))

有人可以解释为什么会发生这种情况吗？

编辑:

这发生在 0.23.1 上，但没有发生在 0.23.1 上，所以这可能是一个错误。

Answer 1

我认为你需要 .name 属性:

df1 = df.groupby('y').apply(lambda x: x.sample(sizes[x.name]))
print (df1)

     y  x  z
y           
0 4  0  1  2
  2  0  1  2
1 6  1  1  2
2 0  2  1  2

如果可能某些值在字典中不匹配，请使用 get 和 0 来获取不匹配的值:

df1 = df.groupby('y').apply(lambda x: x.sample(sizes.get(x.name, 0)))

编辑:

问题是唯一返回一个元素numpy数组:

def f(x):
    print (x['y'].unique())
    print (x['y'].unique()[0])
    print (sizes[x['y'].unique()[0]])
    print (x.sample(sizes[x['y'].unique()[0]]))

df1 = df.groupby('y').apply(f)

[0]
0
2
   y  x  z
2  0  1  2
4  0  1  2
[0]
0
2
   y  x  z
4  0  1  2
2  0  1  2
[1]
1
1
   y  x  z
6  1  1  2
[2]
2
1
   y  x  z
0  2  1  2

<小时/>

df1 = df.groupby('y').apply(lambda x: x.sample(sizes[x.y.unique()[0]]))
print (df1)
     y  x  z
y           
0 4  0  1  2
  2  0  1  2
1 6  1  1  2
2 0  2  1  2