当前位置：首页 > news >正文

别再手动分数据集了！用Python实现KS算法自动划分训练集和测试集（附完整代码）

news 2026/5/1 21:54:01

用KS算法智能划分数据集：告别随机分割的五大痛点

在机器学习的第一个实战环节——数据准备阶段，许多初学者都会不假思索地使用train_test_split进行随机划分。这种看似便捷的操作，却可能为后续模型评估埋下隐患。想象一下这样的场景：你的光谱数据中某个关键浓度区间的样本全部被分到了测试集，而训练集恰好缺失了这个特征区间的代表性数据，最终模型在实际预测时就会出现系统性偏差。

Kennard-Stone算法（简称KS算法）正是为解决这类问题而生。这个诞生于化学计量学领域的经典算法，通过最大化样本在特征空间中的覆盖度，确保训练集和测试集都能完整反映原始数据的分布特性。与随机划分相比，它能有效避免数据划分偏差导致的模型评估失真，特别适用于光谱数据、分子描述符等高维特征的数据集划分。

1. 为什么需要KS算法：随机划分的五大缺陷

随机划分数据集就像蒙着眼睛玩飞镖——结果完全不可控。以下是实践中常见的五大问题：

关键特征区间遗漏：某些特征值区间可能完全被排除在训练集外
离群点分配不均：重要但稀有的离群点可能全部进入测试集
类别不平衡加剧：随机可能放大原始数据中的类别不平衡问题
时间序列断裂：对时间相关数据会破坏序列连续性
评估结果波动大：不同随机种子导致评估指标差异显著

# 随机划分的典型问题演示 from sklearn.model_selection import train_test_split import numpy as np # 模拟存在明显聚类特征的数据 np.random.seed(42) cluster1 = np.random.normal(loc=0, scale=0.5, size=(50, 2)) cluster2 = np.random.normal(loc=5, scale=1, size=(30, 2)) X = np.vstack([cluster1, cluster2]) y = np.array([0]*50 + [1]*30) # 随机划分可能导致的极端情况 X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3) print("训练集类别分布:", np.bincount(y_train)) print("测试集类别分布:", np.bincount(y_test))

执行上述代码，你可能会发现测试集中完全缺失某个类别的情况。而KS算法通过欧氏距离的全局优化，从根本上杜绝了这类问题。

2. KS算法核心原理与数学实现

KS算法的智慧在于它模拟了人类选择代表性样本的思维方式：先确定边界点，再逐步填充内部。其核心步骤可分解为：

寻找初始锚点：选择欧氏距离最远的两个样本作为初始训练集
迭代扩展：每次选择距离已选样本最远的点加入训练集
终止条件：当训练集数量达到预设比例时停止

算法依赖的欧氏距离计算公式为：

$$ d(\mathbf{x}p, \mathbf{x}q) = \sqrt{\sum{i=1}^n (x{p,i} - x_{q,i})^2} $$

其中$\mathbf{x}_p$和$\mathbf{x}_q$代表两个样本向量，$n$为特征维度。

import numpy as np def ks_train_test_split(X, y, test_size=0.2): """ KS算法实现的数据集划分 参数: X: 特征矩阵 (n_samples, n_features) y: 标签向量 (n_samples,) test_size: 测试集比例 返回: 划分后的训练集和测试集 """ n_samples = X.shape[0] train_size = int(n_samples * (1 - test_size)) selected_indices = [] # 步骤1：选择距离均值最远的样本作为第一个点 distances_to_mean = np.sum((X - np.mean(X, axis=0))**2, axis=1) first_point = np.argmax(distances_to_mean) selected_indices.append(first_point) # 步骤2：选择距离第一个点最远的样本作为第二个点 distances = np.sum((X - X[first_point])**2, axis=1) second_point = np.argmax(distances) selected_indices.append(second_point) # 步骤3：迭代选择剩余点 while len(selected_indices) < train_size: # 计算每个未选点到已选点的最小距离 min_distances = [] for i in range(n_samples): if i not in selected_indices: dist = np.min(np.sum((X[selected_indices] - X[i])**2, axis=1)) min_distances.append((i, dist)) # 选择最小距离最大的样本 min_distances.sort(key=lambda x: x[1], reverse=True) next_point = min_distances[0][0] selected_indices.append(next_point) # 生成测试集索引 all_indices = set(range(n_samples)) test_indices = list(all_indices - set(selected_indices)) return X[selected_indices], X[test_indices], y[selected_indices], y[test_indices]

3. 实战对比：KS算法 vs 随机划分

为了直观展示KS算法的优势，我们用一个实际的光谱数据集进行对比实验。使用scikit-learn中的葡萄酒数据集作为示例：

from sklearn.datasets import load_wine import matplotlib.pyplot as plt from sklearn.manifold import TSNE # 加载数据 data = load_wine() X, y = data.data, data.target # 两种划分方式 X_train_rand, X_test_rand, y_train_rand, y_test_rand = train_test_split( X, y, test_size=0.3, random_state=42) X_train_ks, X_test_ks, y_train_ks, y_test_ks = ks_train_test_split( X, y, test_size=0.3) # t-SNE可视化 def plot_tsne(X_train, X_test, title): tsne = TSNE(n_components=2, random_state=42) X_combined = np.vstack([X_train, X_test]) X_tsne = tsne.fit_transform(X_combined) plt.figure(figsize=(12, 5)) plt.scatter(X_tsne[:len(X_train), 0], X_tsne[:len(X_train), 1], c='blue', label='Train set', alpha=0.6) plt.scatter(X_tsne[len(X_train):, 0], X_tsne[len(X_train):, 1], c='red', label='Test set', alpha=0.6) plt.title(title) plt.legend() plt.show() plot_tsne(X_train_rand, X_test_rand, "Random Split") plot_tsne(X_train_ks, X_test_ks, "KS Algorithm Split")

通过可视化对比可以明显看出，随机划分的训练集（蓝色）在某些区域过度集中，而测试集（红色）则在其他区域形成孤岛。KS算法的划分结果则呈现出更均匀的空间覆盖，确保了两个集合都能全面代表原始数据分布。

4. 高级应用技巧与性能优化

基础版的KS算法虽然有效，但在处理大规模数据时可能面临计算效率问题。以下是几个提升实用性的进阶技巧：

4.1 距离计算优化

原始算法需要计算所有样本间的距离矩阵，时间复杂度为O(n²)。可以采用以下优化策略：

# 使用矩阵运算加速距离计算 def optimized_distance(X): X_sq = np.sum(X**2, axis=1) dists = X_sq[:, None] + X_sq[None, :] - 2 * X @ X.T np.fill_diagonal(dists, 0) return np.sqrt(dists) # 分块处理大数据集 def chunked_ks(X, y, test_size=0.2, chunk_size=1000): n_samples = X.shape[0] if n_samples <= chunk_size: return ks_train_test_split(X, y, test_size) # 先对数据聚类，再从每个簇中按比例选取 from sklearn.cluster import KMeans kmeans = KMeans(n_clusters=chunk_size) clusters = kmeans.fit_predict(X) selected_indices = [] for cluster_id in np.unique(clusters): mask = clusters == cluster_id X_cluster = X[mask] y_cluster = y[mask] _, _, _, _ = ks_train_test_split( X_cluster, y_cluster, test_size=test_size) # 省略具体实现细节...

4.2 与分层抽样的结合

对于类别不平衡数据，可以将KS算法与分层抽样结合：

def stratified_ks(X, y, test_size=0.2): unique_classes = np.unique(y) train_indices = [] test_indices = [] for cls in unique_classes: mask = y == cls X_cls = X[mask] y_cls = y[mask] n_cls_samples = sum(mask) n_test = max(1, int(n_cls_samples * test_size)) # 对每个类别单独应用KS算法 X_train_cls, X_test_cls, y_train_cls, y_test_cls = ks_train_test_split( X_cls, y_cls, test_size=n_test/n_cls_samples) # 收集索引 original_indices = np.where(mask)[0] train_indices.extend(original_indices[y_train_cls == cls]) test_indices.extend(original_indices[y_test_cls == cls]) return X[train_indices], X[test_indices], y[train_indices], y[test_indices]

4.3 自定义距离度量

欧氏距离并非适用于所有场景。对于文本或基因序列等特殊数据，可以替换为其他距离度量：

from scipy.spatial.distance import cosine def ks_with_custom_metric(X, y, test_size=0.2, metric='euclidean'): if metric == 'cosine': def distance_func(a, b): return cosine(a, b) else: def distance_func(a, b): return np.linalg.norm(a - b) # 修改原始KS算法中的距离计算部分 # 具体实现省略...

5. 工程实践中的注意事项

在实际项目中应用KS算法时，有几个关键点需要特别注意：

特征缩放的重要性：KS算法基于距离度量，不同特征尺度会影响结果

from sklearn.preprocessing import StandardScaler scaler = StandardScaler() X_scaled = scaler.fit_transform(X)

分类与回归任务的差异：
- 分类任务：建议先按类别分层，再在每个类别内应用KS算法
- 回归任务：可直接应用，或对目标变量分箱后分层处理

计算资源管理：

对于>10万样本的数据，考虑使用近似算法或基于聚类的分块策略

可将距离矩阵计算转移到GPU加速：

import cupy as cp def gpu_distance(X): X_gpu = cp.array(X) dists = cp.sum(X_gpu**2, axis=1)[:, None] + \ cp.sum(X_gpu**2, axis=1)[None, :] - \ 2 * X_gpu @ X_gpu.T return cp.asnumpy(dists)

与交叉验证的配合：

from sklearn.model_selection import KFold # 先使用KS算法划分出干净的测试集 X_train, X_test, y_train, y_test = ks_train_test_split(X, y, test_size=0.2) # 在训练集上应用交叉验证 kf = KFold(n_splits=5) for train_index, val_index in kf.split(X_train): X_tr, X_val = X_train[train_index], X_train[val_index] y_tr, y_val = y_train[train_index], y_train[val_index] # 训练和验证...