pyhton中K折交叉验证StratifiedKFold_stratifiedkfoldpython

https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.StratifiedKFold.html#sklearn.model_selection.StratifiedKFold

from sklearn.model_selection import StratifiedKFold
X = np.array([[1, 2], [3, 4], [1, 2], [3, 4], [1, 2], [3, 4], [1, 2], [3, 4], [1, 2], [3, 4], [1, 2], [3, 4]])
y = np.array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2])
#y = np.hstack(([0]*4,[1]*4,[2]*4))
skf = StratifiedKFold(n_splits=3, random_state=None, shuffle=True)
skf.get_n_splits(X, y)
 
print(skf)  
 
for train_index, test_index in skf.split(X, y):
    print("TRAIN:", train_index, "TEST:", test_index)
    X_train, X_test = X[train_index], X[test_index]
    y_train, y_test = y[train_index], y[test_index]
 
StratifiedKFold(n_splits=3, random_state=None, shuffle=True)
TRAIN: [ 1  3  6  7  8 10] TEST: [ 0  2  4  5  9 11]
TRAIN: [ 0  1  2  4  5  7  8  9 11] TEST: [ 3  6 10]
TRAIN: [ 0  2  3  4  5  6  9 10 11] TEST: [1 7 8]