K近邻(K-Nearest Neighbors)是一种直观且强大的非参数分类与回归算法。它不需要显式的训练过程,而是通过在特征空间中寻找最近的K个邻居来进行预测。本文用Java完整实现KNN算法,讲解多种距离度量、K值选择与KD树加速优化,帮助读者掌握这一机器学习入门经典。
一、算法原理
KNN的核心思想非常朴素:”近朱者赤,近墨者黑”。给定一个待分类样本,在训练集中找到与它距离最近的K个样本,然后根据这K个邻居的类别进行投票,得票最多的类别即为预测结果。
算法的关键要素包括:
- 距离度量:如何定义样本之间的”远近”;
- K值选择:邻居数量对模型偏差与方差的影响;
- 投票策略:简单多数投票或加权投票。
二、数据结构与距离度量
import java.util.*;
/**
* KNN算法完整实现
* 支持多种距离度量与KD树加速
*/
public class KNNClassifier {
/**
* 样本点:包含特征向量和标签
*/
public static class Sample {
double[] features;
String label;
public Sample(double[] features, String label) {
this.features = features.clone();
this.label = label;
}
@Override
public String toString() {
return Arrays.toString(features) + " -> " + label;
}
}
/**
* 邻居结果
*/
public static class Neighbor implements Comparable<Neighbor> {
Sample sample;
double distance;
public Neighbor(Sample sample, double distance) {
this.sample = sample;
this.distance = distance;
}
@Override
public int compareTo(Neighbor other) {
return Double.compare(this.distance, other.distance);
}
}
// 训练数据集
private List<Sample> trainingData;
// K值
private int k;
// 距离度量类型
private DistanceMetric metric;
// 是否使用加权投票
private boolean weighted;
/**
* 距离度量接口
*/
public interface DistanceMetric {
double distance(double[] a, double[] b);
}
/**
* 欧氏距离
*/
public static final DistanceMetric EUCLIDEAN = (a, b) -> {
double sum = 0;
for (int i = 0; i < a.length; i++) {
double diff = a[i] - b[i];
sum += diff * diff;
}
return Math.sqrt(sum);
};
/**
* 曼哈顿距离
*/
public static final DistanceMetric MANHATTAN = (a, b) -> {
double sum = 0;
for (int i = 0; i < a.length; i++) {
sum += Math.abs(a[i] - b[i]);
}
return sum;
};
/**
* 切比雪夫距离
*/
public static final DistanceMetric CHEBYSHEV = (a, b) -> {
double max = 0;
for (int i = 0; i < a.length; i++) {
max = Math.max(max, Math.abs(a[i] - b[i]));
}
return max;
};
/**
* 闵可夫斯基距离(p=3)
*/
public static final DistanceMetric MINKOWSKI3 = (a, b) -> {
double sum = 0;
for (int i = 0; i < a.length; i++) {
sum += Math.pow(Math.abs(a[i] - b[i]), 3);
}
return Math.cbrt(sum);
};
public KNNClassifier(int k, DistanceMetric metric, boolean weighted) {
this.k = k;
this.metric = metric;
this.weighted = weighted;
this.trainingData = new ArrayList<>();
}
/**
* 训练:KNN是惰性学习,只需存储训练数据
*/
public void train(List<Sample> data) {
this.trainingData = new ArrayList<>(data);
}
/**
* 预测单个样本的类别
*/
public String predict(double[] features) {
// 计算与所有训练样本的距离
PriorityQueue<Neighbor> maxHeap = new PriorityQueue<>(Collections.reverseOrder());
for (Sample sample : trainingData) {
double dist = metric.distance(features, sample.features);
maxHeap.offer(new Neighbor(sample, dist));
// 维护大小为K的最大堆
if (maxHeap.size() > k) {
maxHeap.poll();
}
}
// 收集K个最近邻居
List<Neighbor> neighbors = new ArrayList<>(maxHeap);
// 投票
Map<String, Double> votes = new HashMap<>();
for (Neighbor n : neighbors) {
String label = n.sample.label;
double weight = weighted ? 1.0 / (n.distance + 1e-10) : 1.0;
votes.merge(label, weight, Double::sum);
}
// 返回得票最高的类别
String bestLabel = null;
double maxVote = -1;
for (Map.Entry<String, Double> entry : votes.entrySet()) {
if (entry.getValue() > maxVote) {
maxVote = entry.getValue();
bestLabel = entry.getKey();
}
}
return bestLabel;
}
/**
* 批量预测
*/
public List<String> predictBatch(List<double[]> featuresList) {
List<String> results = new ArrayList<>();
for (double[] features : featuresList) {
results.add(predict(features));
}
return results;
}
三、KD树加速
当训练数据量很大时,线性扫描所有样本的效率很低。KD树通过递归地将空间划分为超矩形区域,将最近邻搜索的复杂度从O(N)降低到O(log N)。
/**
* KD树节点
*/
private static class KDNode {
Sample point;
int splitDim; // 分裂维度
KDNode left, right; // 左右子树
KDNode(Sample point, int splitDim) {
this.point = point;
this.splitDim = splitDim;
}
}
// KD树根节点
private KDNode kdRoot;
/**
* 构建KD树
*/
public void buildKDTree() {
if (trainingData.isEmpty()) return;
kdRoot = buildKDTreeRecursive(new ArrayList<>(trainingData), 0);
}
private KDNode buildKDTreeRecursive(List<Sample> points, int depth) {
if (points.isEmpty()) return null;
int dim = points.get(0).features.length;
int axis = depth % dim;
// 按当前维度排序
points.sort(Comparator.comparingDouble(p -> p.features[axis]));
int median = points.size() / 2;
KDNode node = new KDNode(points.get(median), axis);
node.left = buildKDTreeRecursive(points.subList(0, median), depth + 1);
node.right = buildKDTreeRecursive(points.subList(median + 1, points.size()), depth + 1);
return node;
}
/**
* 使用KD树进行K近邻搜索
*/
public String predictKDTree(double[] features) {
if (kdRoot == null) {
buildKDTree();
}
PriorityQueue<Neighbor> maxHeap = new PriorityQueue<>(Collections.reverseOrder());
searchKDTree(kdRoot, features, maxHeap);
// 投票
Map<String, Double> votes = new HashMap<>();
List<Neighbor> neighbors = new ArrayList<>(maxHeap);
for (Neighbor n : neighbors) {
String label = n.sample.label;
double weight = weighted ? 1.0 / (n.distance + 1e-10) : 1.0;
votes.merge(label, weight, Double::sum);
}
String bestLabel = null;
double maxVote = -1;
for (Map.Entry<String, Double> entry : votes.entrySet()) {
if (entry.getValue() > maxVote) {
maxVote = entry.getValue();
bestLabel = entry.getKey();
}
}
return bestLabel;
}
private void searchKDTree(KDNode node, double[] target, PriorityQueue<Neighbor> maxHeap) {
if (node == null) return;
// 计算当前节点与目标点的距离
double dist = metric.distance(target, node.point.features);
maxHeap.offer(new Neighbor(node.point, dist));
if (maxHeap.size() > k) {
maxHeap.poll();
}
// 决定先搜索哪一侧
int dim = node.splitDim;
KDNode first = target[dim] < node.point.features[dim] ? node.left : node.right;
KDNode second = target[dim] < node.point.features[dim] ? node.right : node.left;
// 递归搜索优先侧
searchKDTree(first, target, maxHeap);
// 检查是否需要搜索另一侧(超球体与分裂超平面相交)
if (maxHeap.size() < k || Math.abs(target[dim] - node.point.features[dim]) < maxHeap.peek().distance) {
searchKDTree(second, target, maxHeap);
}
}
四、交叉验证与K值选择
/**
* K折交叉验证选择最优K值
* @param data 完整数据集
* @param maxK 测试的最大K值
* @param folds 折数
* @return 最优K值
*/
public static int chooseBestK(List<Sample> data, int maxK, int folds, DistanceMetric metric) {
int bestK = 1;
double bestAccuracy = 0;
Collections.shuffle(data, new Random(42));
for (int testK = 1; testK <= maxK; testK += 2) {
double totalAccuracy = 0;
for (int fold = 0; fold < folds; fold++) {
List<Sample> trainSet = new ArrayList<>();
List<Sample> testSet = new ArrayList<>();
for (int i = 0; i < data.size(); i++) {
if (i % folds == fold) {
testSet.add(data.get(i));
} else {
trainSet.add(data.get(i));
}
}
KNNClassifier knn = new KNNClassifier(testK, metric, false);
knn.train(trainSet);
int correct = 0;
for (Sample s : testSet) {
String pred = knn.predict(s.features);
if (pred.equals(s.label)) {
correct++;
}
}
totalAccuracy += (double) correct / testSet.size();
}
double avgAccuracy = totalAccuracy / folds;
System.out.printf("K=%d, 平均准确率: %.4f%n", testK, avgAccuracy);
if (avgAccuracy > bestAccuracy) {
bestAccuracy = avgAccuracy;
bestK = testK;
}
}
System.out.printf("最优K值: %d (准确率: %.4f)%n", bestK, bestAccuracy);
return bestK;
}
/**
* 评估模型准确率
*/
public double evaluate(List<Sample> testData) {
int correct = 0;
for (Sample s : testData) {
String pred = predict(s.features);
if (pred.equals(s.label)) {
correct++;
}
}
return (double) correct / testData.size();
}
/**
* 打印混淆矩阵
*/
public void printConfusionMatrix(List<Sample> testData) {
Set<String> labels = new TreeSet<>();
for (Sample s : testData) labels.add(s.label);
List<String> labelList = new ArrayList<>(labels);
int[][] matrix = new int[labelList.size()][labelList.size()];
for (Sample s : testData) {
String pred = predict(s.features);
int actualIdx = labelList.indexOf(s.label);
int predIdx = labelList.indexOf(pred);
matrix[actualIdx][predIdx]++;
}
System.out.println("混淆矩阵:");
System.out.print(" ");
for (String l : labelList) System.out.printf("%6s", l);
System.out.println();
for (int i = 0; i < labelList.size(); i++) {
System.out.printf("%6s", labelList.get(i));
for (int j = 0; j < labelList.size(); j++) {
System.out.printf("%6d", matrix[i][j]);
}
System.out.println();
}
}
五、完整测试示例
public static void main(String[] args) {
// 构造鸢尾花简化数据集(2个特征:花萼长度、花萼宽度)
List<Sample> dataset = new ArrayList<>();
Random rand = new Random(42);
// 类别A:特征较小
for (int i = 0; i < 50; i++) {
dataset.add(new Sample(new double[]{
5.0 + rand.nextGaussian() * 0.5,
3.0 + rand.nextGaussian() * 0.3
}, "A"));
}
// 类别B:特征中等
for (int i = 0; i < 50; i++) {
dataset.add(new Sample(new double[]{
6.0 + rand.nextGaussian() * 0.5,
3.5 + rand.nextGaussian() * 0.3
}, "B"));
}
// 类别C:特征较大
for (int i = 0; i < 50; i++) {
dataset.add(new Sample(new double[]{
7.0 + rand.nextGaussian() * 0.5,
4.0 + rand.nextGaussian() * 0.3
}, "C"));
}
// 划分训练集和测试集
Collections.shuffle(dataset, rand);
int trainSize = (int)(dataset.size() * 0.8);
List<Sample> trainSet = dataset.subList(0, trainSize);
List<Sample> testSet = dataset.subList(trainSize, dataset.size());
System.out.println("=== KNN分类器测试 ===");
System.out.printf("训练集大小: %d, 测试集大小: %d%n%n", trainSet.size(), testSet.size());
// 1. 选择最优K值
System.out.println("选择最优K值(5折交叉验证):");
int bestK = chooseBestK(new ArrayList<>(dataset), 15, 5, EUCLIDEAN);
// 2. 使用最优K值训练并评估
System.out.println("\n使用最优K值评估:");
KNNClassifier knn = new KNNClassifier(bestK, EUCLIDEAN, true);
knn.train(trainSet);
double accuracy = knn.evaluate(testSet);
System.out.printf("测试集准确率: %.4f%n%n", accuracy);
// 3. 比较不同距离度量
System.out.println("不同距离度量比较:");
DistanceMetric[] metrics = {EUCLIDEAN, MANHATTAN, CHEBYSHEV, MINKOWSKI3};
String[] metricNames = {"欧氏距离", "曼哈顿距离", "切比雪夫距离", "闵可夫斯基(p=3)"};
for (int i = 0; i < metrics.length; i++) {
KNNClassifier testKnn = new KNNClassifier(bestK, metrics[i], true);
testKnn.train(trainSet);
double acc = testKnn.evaluate(testSet);
System.out.printf("%s: %.4f%n", metricNames[i], acc);
}
// 4. KD树加速效果对比
System.out.println("\nKD树加速效果:");
KNNClassifier kdKnn = new KNNClassifier(bestK, EUCLIDEAN, true);
kdKnn.train(trainSet);
kdKnn.buildKDTree();
long start1 = System.currentTimeMillis();
for (Sample s : testSet) {
knn.predict(s.features);
}
long linearTime = System.currentTimeMillis() - start1;
long start2 = System.currentTimeMillis();
for (Sample s : testSet) {
kdKnn.predictKDTree(s.features);
}
long kdTime = System.currentTimeMillis() - start2;
System.out.printf("线性扫描耗时: %dms%n", linearTime);
System.out.printf("KD树搜索耗时: %dms%n", kdTime);
// 5. 打印混淆矩阵
System.out.println("\n混淆矩阵(线性扫描):");
knn.printConfusionMatrix(testSet);
}
}
六、复杂度分析
- 线性扫描预测:时间复杂度O(N × d),N为训练样本数,d为特征维度;空间复杂度O(N × d)存储训练数据。
- KD树构建:时间复杂度O(N log N),递归排序中位数;空间复杂度O(N)。
- KD树搜索:平均时间复杂度O(log N),最坏情况下(所有点在同一区域)退化为O(N)。
- 交叉验证:K折交叉验证时间复杂度为O(K × N² × d),适合离线调参。
七、总结
本文完整实现了KNN算法及其优化版本,核心内容包括:
- 多种距离度量:欧氏距离、曼哈顿距离、切比雪夫距离、闵可夫斯基距离,适应不同数据分布;
- 加权投票策略:距离越近的邻居权重越高,降低噪声样本的影响;
- KD树加速:通过空间划分将最近邻搜索从线性降至对数级别,大幅提升大数据集上的预测效率;
- 交叉验证选参:使用K折交叉验证自动选择最优K值,避免过拟合与欠拟合。
KNN虽然简单,但在许多实际场景中表现优异,尤其适用于数据集规模适中、决策边界复杂的分类任务。理解KNN也是后续学习核方法、度量学习等高级主题的重要基础。