How do I do stratified sampling on group-separated datasets in Python? Do packages for this exist? Say I have the following data:
    Group_ID | Column_1 | Column_2 | Column_3 ...
==========================================
A        | 1        | 2        | 33
A        | 2        | 2        | 3765
A        | 3        | 6        | 3436
A        | 4        | 8        | 32
B        | 5        | 9        | 33
B        | 3        | 34       | 385
B        | 7        | 25       | 3
B        | 3        | 1        | 38
C        | 6        | 2        | 3
C        | 8        | 2        | 4
D        | 7        | 1        | 5
D        | 6        | 9        | 11

I want to:


*

*First identify train-test splits that keep groups (Group_ID) separate between splits. I.e. no group can be in both train and test splits.

*Out of all possible splits that have been identified, get splits which have the most similar distributions of Column_1, Column_2, Column_3 etc. across train and test splits.


In short, is there any way that I can split my data so that groups are separated, but that the other features are similar across the split?
Ideally, I would like to do this with a package in Python or the like, if it exists.
 A: This problem is much more complicated than it may look. Try to construct a solution for your simple example. In general, this problem in general can be interpreted as multidimensional knapsack problem and thus even approximations are very hard to find - the one-dimensional knapsack problem is already NP-hard. In your toy example one can simply try all possibilities and choose the 'best' one. Which brings to the second problem: What is the 'best' solution? The goal is to have a proper split for all classes and thus this is a multi-critera optimization problem, so there may be solutions that are clearly worse than others, but one solution may be better for one class and a second solution may be better for another class. One possibility might consider the hypervolume created by the solutions with respect to the theoretical optimal solution. The theoretical optimal solution would create two sets such that both contain half of the samples of each class. For your toy example, I think the optimal solution (with respect to the product of the hypervolume of both splits) is to use: 
split1 = Group A + Group B, and split2 = Group C + Group D
This yields: 
| split1 | Column_1 | Column_2 | Column 3 | 
================================== 
A+B    | 28       | 87       | 7725 
C+D    | 27       | 14       | 23 

Column 1 is split really good: 27.5 would be optimal and it reaches 27 and 28. 
Column 2 is split rather bad: optimal would be 50.5, and we get 87 and 14. 
Column 3 is split really bad: optimal would be 3874, and we get 7725 and 23. 
Therefore, I assume that there is no package to solve your problem. Good luck anyway!
