# Stratified Kfold splits by inputs instead of just being based on outputs

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])