基于字典的样本 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]]))

y y x z

0 2 0 1 2 4 0 1 2 1 5 1 1 2 2 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)))

编辑：

问题是 unique return 一个元素 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

基于字典的样本 Pandas

Sample Pandas based on dictionary

python

sample

pandas

keyerror