nykergoto’s blog

機械学習とpythonをメインに

連続変数で StratifiedKFold

pandas.qcut で int に変換すると良さげ

from sklearn.model_selection import StratifiedKFold
import pandas as pd

def get_stratified_fold_split(X, y, n_folds=10, q=20):
    # Categorical.codes で各クラスに対応する int が取得できる
    y_cat = pd.qcut(y, q=q).codes
    fold = StratifiedKFold(n_splits=n_folds, random_state=71)
    return fold.split(X, y_cat)

qint 化するときの分割数なので n_folds に合わせていい感じに合わせると良いかな?