每日算法 — 使用java实现KNN:K近邻分类与距离度量优化

K近邻(K-Nearest Neighbors)是一种直观且强大的非参数分类与回归算法。它不需要显式的训练过程,而是通过在特征空间中寻找最近的K个邻居来进行预测。本文用Java完整实现KNN算法,讲解多种距离度量、K值选择与KD树加速优化,帮助读者掌握这一机器学习入门经典。

一、算法原理

KNN的核心思想非常朴素:”近朱者赤,近墨者黑”。给定一个待分类样本,在训练集中找到与它距离最近的K个样本,然后根据这K个邻居的类别进行投票,得票最多的类别即为预测结果。

算法的关键要素包括:

  1. 距离度量:如何定义样本之间的”远近”;
  2. K值选择:邻居数量对模型偏差与方差的影响;
  3. 投票策略:简单多数投票或加权投票。

二、数据结构与距离度量

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算法及其优化版本,核心内容包括:

  1. 多种距离度量:欧氏距离、曼哈顿距离、切比雪夫距离、闵可夫斯基距离,适应不同数据分布;
  2. 加权投票策略:距离越近的邻居权重越高,降低噪声样本的影响;
  3. KD树加速:通过空间划分将最近邻搜索从线性降至对数级别,大幅提升大数据集上的预测效率;
  4. 交叉验证选参:使用K折交叉验证自动选择最优K值,避免过拟合与欠拟合。

KNN虽然简单,但在许多实际场景中表现优异,尤其适用于数据集规模适中、决策边界复杂的分类任务。理解KNN也是后续学习核方法、度量学习等高级主题的重要基础。