基于字典的 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 返回一个元素的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