K-fold cross-validation splits

Write codeMedium25 minMeta style

Your task

Generate k-fold cross-validation index splits for a dataset of n samples (indices 0..n-1, in order, no shuffling). Split the indices into k contiguous folds as evenly as possible: if n is not divisible by k, the first (n mod k) folds get one extra element. For each fold i, that fold is the validation set and the remaining indices (in ascending order) form the training set. Return a list of k pairs [train_indices, val_indices]. Assume 1 <= k <= n.

Examples

kfold_indices(n: int, k: int) -> list[list[list[int]]]
Example 1
Inputkfold_indices(6, 3)
Output[[[2, 3, 4, 5], [0, 1]], [[0, 1, 4, 5], [2, 3]], [[0, 1, 2, 3], [4, 5]]]