python中带约束的分层抽样

Question

我有一个观察数据框

data = [['red', 1, 0.2], ['blue', 1, 0.5], ['green', 2, 0.8], ['blue', 2, 0.55], ['blue', 2, 0.52], ['red', 3, 0.15], ['green', 3, 0.85], ['red', 4, 0.12], ['purple', 4, 0.01]] 

df = pd.DataFrame(data, columns = ['label', 'signal', 'value']) 

   label    signal  value
0   red     1   0.20
1   blue    1   0.50
2   green   2   0.80
3   blue    2   0.55
4   blue    2   0.52
5   red     3   0.15
6   green   3   0.85
7   red     4   0.12
8   purple  4   0.01

我想在标签上做分层k折叠采样，但是我需要这样做，这样就没有信号值在折叠中被分割。我用一个只使用字典和复杂检查的实现来实现它。我想知道有没有更简单的方法来解决这个问题？你知道吗

K=2的结果可能是：

batch 1
0   red     1   0.20
1   blue    1   0.50
5   red     3   0.15
6   green   3   0.85

batch 2
2  green    2   0.80
3   blue    2   0.55
4   blue    2   0.52
7   red     4   0.12
8   purple  4   0.01

其中第1批有2红1蓝1绿，第2批有1红2蓝1绿1紫。在这种情况下，两个批次在类内容方面是平衡的，这就是我想要的。你知道吗

Answer 1

我认为您正在寻找内置于scikit learn中的GroupShuffleSplit函数： sklearn.model_selection.GroupShuffleSplit