K-Fold Cross Validation in Python: A Simple Implementation

import random
from statistics import mean


def cross_validation_scores(data, labels, k=5, seed=42):
    random.seed(seed)
    indices = list(range(len(data)))
    random.shuffle(indices)
    fold_size = len(indices) // k
    folds = []
    for i in range(k):
        if i == k - 1:
            folds.append(indices[i * fold_size:])
        else:
            folds.append(indices[i * fold_size:(i + 1) * fold_size])

    scores = []
    for i in range(k):
        test_idx = folds[i]
        train_idx = [idx for j in range(k) if j != i for idx in folds[j]]
        train_labels = [labels[idx] for idx in train_idx]
        test_labels = [labels[idx] for idx in test_idx]
        train_mean = mean(train_labels)
        predictions = [train_mean] * len(test_labels)
        mse = mean((pred - actual) ** 2 for pred, actual in zip(predictions, test_labels))
        scores.append(mse)
    return scores


if __name__ == "__main__":
    data = list(range(20))
    labels = [x * 0.5 + 1 for x in data]
    mse_scores = cross_validation_scores(data, labels, k=5)
    print("MSE scores per fold:", mse_scores)
    print("Average MSE:", mean(mse_scores))