0
$\begingroup$

It's known that a stratified split can easily be done if there is a known imbalance in the output distribution (or output classes).

Is it possible to do a stratified k-fold using the input distribution instead of just with the output distribution?

For instance, say I'm doing regression on a dataset with

(a) a single output variable that has a fairly uniform distribution and

(b) an input feature with 3 distinct values where one values occurs in 50% of the samples and the other 2 values occur 25% each. I'm trying to create 2 splits for them.

from sklearn.model_selection import StratifiedKFold

X = np.array([20, 20, 20, 20, 30, 30, 40, 40])
y = np.ones(8)
skf = StratifiedKFold(n_splits=2)
for a, b in skf.split(X, y):
    print('train: ', a, X[a], 'test: ', b , X[b])

# actual answer
train:  [4 5 6 7] [30 30 40 40] test:  [0 1 2 3] [20 20 20 20]
train:  [0 1 2 3] [20 20 20 20] test:  [4 5 6 7] [30 30 40 40]

# answer I hoped to see
train:  [0 1 4 6] [20 20 30 40] test:  [2 3 5 7] [20 20 30 40]
train:  [2 3 5 7] [20 20 30 40] test:  [0 1 4 6] [20 20 30 40]

Is it possible to do the stratified folds based on the input distribution? My actual problem has ~10 columns for X where 9 columns have a uniform distribution and one column has this skewed sort of distribution.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.