每日算法 — 使用java实现k-d树：多维空间最近邻搜索与范围查询

可莉爸爸 AI算法

算法概述

k-d树（k-dimensional tree）是由Jon Bentley于1975年提出的空间划分数据结构，通过递归地沿某一维度的中位数值将k维空间二分，构建一棵平衡二叉搜索树。与四叉树的固定四等分不同，k-d树按维度循环选择划分轴，能够优雅地扩展到任意维度空间。

k-d树在高维最近邻搜索、碰撞检测、机器学习中的KNN算法、图像特征匹配等领域有着不可替代的作用。本文用Java实现一棵完整的k-d树，支持动态插入、最近邻搜索、K近邻搜索和正交范围查询。

核心原理

k-d树的构建遵循以下规则：
1. 选择划分维度：在深度 (d) 处，选择维度 (d \mod k) 作为划分轴
2. 选择划分点：取该维度上的中位数作为节点值，将剩余点分为左右两子树
3. 递归构建：对左右子集分别递归执行上述过程

划分超平面垂直于当前维度，将空间一分为二。搜索时利用超平面的几何性质进行剪枝，避免遍历整棵树。

Java实现

项目结构

kdtree/
├── KDPoint.java           // k维点封装
├── KDNode.java            // k-d树节点
├── KDTree.java            // k-d树核心实现
├── KNNClassifier.java     // KNN分类器应用
└── Main.java              // 测试入口

基础数据结构

import java.util.Arrays;

/**
 * k维空间点
 */
public class KDPoint {
    public final double[] coords;  // 坐标数组
    public final Object data;      // 附加数据
    public final int dim;          // 维度

    public KDPoint(double[] coords) {
        this(coords, null);
    }

    public KDPoint(double[] coords, Object data) {
        this.coords = Arrays.copyOf(coords, coords.length);
        this.data = data;
        this.dim = coords.length;
    }

    /**
     * 获取指定维度的坐标值
     */
    public double get(int axis) {
        return coords[axis];
    }

    /**
     * 欧几里得距离平方
     */
    public double distSq(KDPoint other) {
        double sum = 0;
        for (int i = 0; i < dim; i++) {
            double diff = coords[i] - other.coords[i];
            sum += diff * diff;
        }
        return sum;
    }

    public double dist(KDPoint other) {
        return Math.sqrt(distSq(other));
    }

    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder("(");
        for (int i = 0; i < dim; i++) {
            sb.append(String.format("%.2f", coords[i]));
            if (i < dim - 1) sb.append(", ");
        }
        sb.append(")");
        return sb.toString();
    }
}

k-d树核心实现：KDTree.java

import java.util.*;

/**
 * k-d树实现
 * 支持构建、插入、最近邻搜索、K近邻搜索、范围查询
 */
public class KDTree {

    private KDNode root;
    private final int k;       // 维度数
    private int size;

    /**
     * k-d树节点
     */
    private static class KDNode {
        KDPoint point;
        int axis;          // 划分维度
        KDNode left, right;

        KDNode(KDPoint point, int axis) {
            this.point = point;
            this.axis = axis;
        }
    }

    /**
     * 按指定维度比较点的比较器
     */
    private static class AxisComparator implements Comparator<KDPoint> {
        private final int axis;

        AxisComparator(int axis) {
            this.axis = axis;
        }

        @Override
        public int compare(KDPoint a, KDPoint b) {
            return Double.compare(a.get(axis), b.get(axis));
        }
    }

    public KDTree(int k) {
        if (k <= 0) throw new IllegalArgumentException("维度必须为正");
        this.k = k;
        this.size = 0;
    }

    /**
     * 从点集构建平衡的k-d树
     * 时间复杂度: O(n log n)
     */
    public void build(List<KDPoint> points) {
        if (points == null || points.isEmpty()) return;
        root = buildRecursive(new ArrayList<>(points), 0);
        size = points.size();
    }

    /**
     * 递归构建
     */
    private KDNode buildRecursive(List<KDPoint> points, int depth) {
        if (points.isEmpty()) return null;

        int axis = depth % k;

        // 按当前维度排序并取中位数
        points.sort(new AxisComparator(axis));
        int medianIdx = points.size() / 2;

        KDNode node = new KDNode(points.get(medianIdx), axis);

        // 递归构建左右子树
        node.left = buildRecursive(points.subList(0, medianIdx), depth + 1);
        node.right = buildRecursive(points.subList(medianIdx + 1, points.size()), depth + 1);

        return node;
    }

    /**
     * 插入单个点
     */
    public void insert(KDPoint point) {
        if (point == null || point.dim != k) {
            throw new IllegalArgumentException("点维度不匹配");
        }
        root = insertRecursive(root, point, 0);
        size++;
    }

    private KDNode insertRecursive(KDNode node, KDPoint point, int depth) {
        if (node == null) {
            return new KDNode(point, depth % k);
        }

        int axis = node.axis;
        if (point.get(axis) < node.point.get(axis)) {
            node.left = insertRecursive(node.left, point, depth + 1);
        } else {
            node.right = insertRecursive(node.right, point, depth + 1);
        }
        return node;
    }

    /**
     * 最近邻搜索
     * @param target 目标点
     * @return 最近的点
     */
    public KDPoint nearest(KDPoint target) {
        if (root == null) return null;

        BestResult best = new BestResult();
        nearestRecursive(root, target, best);
        return best.point;
    }

    private void nearestRecursive(KDNode node, KDPoint target, BestResult best) {
        if (node == null) return;

        // 计算当前节点与目标的距离
        double distSq = node.point.distSq(target);
        if (distSq < best.distSq) {
            best.distSq = distSq;
            best.point = node.point;
        }

        int axis = node.axis;
        double diff = target.get(axis) - node.point.get(axis);

        // 决定优先搜索的子树
        KDNode near = diff < 0 ? node.left : node.right;
        KDNode far = diff < 0 ? node.right : node.left;

        // 先搜索近侧子树
        nearestRecursive(near, target, best);

        // 剪枝：如果目标点到划分超平面的距离已大于当前最优，无需搜索远侧
        if (diff * diff < best.distSq) {
            nearestRecursive(far, target, best);
        }
    }

    private static class BestResult {
        double distSq = Double.POSITIVE_INFINITY;
        KDPoint point = null;
    }

    /**
     * K近邻搜索
     * @param target 目标点
     * @param k_nn   近邻数量
     * @return 按距离排序的K个最近点
     */
    public List<KDPoint> kNearest(KDPoint target, int k_nn) {
        if (root == null || k_nn <= 0) return new ArrayList<>();

        // 最大堆，按距离排序
        PriorityQueue<KDPoint> heap = new PriorityQueue<>(
            (a, b) -> Double.compare(b.distSq(target), a.distSq(target))
        );

        kNearestRecursive(root, target, k_nn, heap);

        List<KDPoint> result = new ArrayList<>(heap);
        result.sort(Comparator.comparingDouble(p -> p.distSq(target)));
        return result;
    }

    private void kNearestRecursive(KDNode node, KDPoint target, int k_nn, 
                                   PriorityQueue<KDPoint> heap) {
        if (node == null) return;

        double distSq = node.point.distSq(target);

        // 如果堆未满，或当前点比堆顶更近，则加入
        if (heap.size() < k_nn) {
            heap.offer(node.point);
        } else if (distSq < heap.peek().distSq(target)) {
            heap.poll();
            heap.offer(node.point);
        }

        int axis = node.axis;
        double diff = target.get(axis) - node.point.get(axis);

        KDNode near = diff < 0 ? node.left : node.right;
        KDNode far = diff < 0 ? node.right : node.left;

        kNearestRecursive(near, target, k_nn, heap);

        // 剪枝条件：远侧超平面距离大于堆中最远点的距离
        double heapMaxDistSq = heap.size() < k_nn ? Double.POSITIVE_INFINITY 
                                                   : heap.peek().distSq(target);
        if (diff * diff < heapMaxDistSq) {
            kNearestRecursive(far, target, k_nn, heap);
        }
    }

    /**
     * 正交范围查询：查找轴对齐矩形区域内的所有点
     */
    public List<KDPoint> rangeQuery(double[] min, double[] max) {
        List<KDPoint> result = new ArrayList<>();
        rangeQueryRecursive(root, min, max, result);
        return result;
    }

    private void rangeQueryRecursive(KDNode node, double[] min, double[] max, 
                                     List<KDPoint> result) {
        if (node == null) return;

        // 检查当前点是否在范围内
        boolean inRange = true;
        for (int i = 0; i < k; i++) {
            if (node.point.get(i) < min[i] || node.point.get(i) > max[i]) {
                inRange = false;
                break;
            }
        }
        if (inRange) {
            result.add(node.point);
        }

        int axis = node.axis;

        // 根据划分决定搜索哪些子树
        if (min[axis] <= node.point.get(axis)) {
            rangeQueryRecursive(node.left, min, max, result);
        }
        if (max[axis] >= node.point.get(axis)) {
            rangeQueryRecursive(node.right, min, max, result);
        }
    }

    public int size() {
        return size;
    }
}

KNN分类器应用：KNNClassifier.java

/**
 * 基于k-d树的K近邻分类器
 * 应用于鸢尾花分类等机器学习任务
 */
public class KNNClassifier {

    public static class Sample {
        KDPoint features;
        String label;

        public Sample(double[] features, String label) {
            this.features = new KDPoint(features, this);
            this.label = label;
        }
    }

    private KDTree tree;
    private final int k;

    public KNNClassifier(int k) {
        this.k = k;
    }

    /**
     * 训练：构建k-d树
     */
    public void train(List<Sample> samples) {
        if (samples == null || samples.isEmpty()) return;

        List<KDPoint> points = new ArrayList<>();
        for (Sample s : samples) {
            points.add(s.features);
        }

        tree = new KDTree(samples.get(0).features.dim);
        tree.build(points);
    }

    /**
     * 预测：查找K近邻并投票
     */
    public String predict(double[] features) {
        if (tree == null) throw new IllegalStateException("未训练");

        KDPoint target = new KDPoint(features);
        List<KDPoint> neighbors = tree.kNearest(target, k);

        // 投票
        Map<String, Integer> votes = new HashMap<>();
        for (KDPoint p : neighbors) {
            Sample s = (Sample) p.data;
            votes.merge(s.label, 1, Integer::sum);
        }

        return votes.entrySet().stream()
            .max(Map.Entry.comparingByValue())
            .map(Map.Entry::getKey)
            .orElse("unknown");
    }

    public static void main(String[] args) {
        // 模拟鸢尾花数据（简化版，仅2维特征）
        List<Sample> trainData = Arrays.asList(
            new Sample(new double[]{5.1, 3.5}, "setosa"),
            new Sample(new double[]{4.9, 3.0}, "setosa"),
            new Sample(new double[]{4.7, 3.2}, "setosa"),
            new Sample(new double[]{4.6, 3.1}, "setosa"),
            new Sample(new double[]{5.0, 3.6}, "setosa"),
            new Sample(new double[]{7.0, 3.2}, "versicolor"),
            new Sample(new double[]{6.4, 3.2}, "versicolor"),
            new Sample(new double[]{6.9, 3.1}, "versicolor"),
            new Sample(new double[]{5.5, 2.3}, "versicolor"),
            new Sample(new double[]{6.5, 2.8}, "versicolor"),
            new Sample(new double[]{6.3, 3.3}, "virginica"),
            new Sample(new double[]{5.8, 2.7}, "virginica"),
            new Sample(new double[]{7.1, 3.0}, "virginica"),
            new Sample(new double[]{6.3, 2.9}, "virginica"),
            new Sample(new double[]{6.5, 3.0}, "virginica")
        );

        KNNClassifier classifier = new KNNClassifier(3);
        classifier.train(trainData);

        // 测试样本
        double[][] tests = {
            {5.0, 3.4},    // 应接近setosa
            {6.7, 3.0},    // 应接近versicolor
            {6.4, 3.1},    // 应接近virginica
        };

        System.out.println("=== KNN分类测试 ===");
        for (double[] test : tests) {
            String pred = classifier.predict(test);
            System.out.printf("特征(%.1f, %.1f) -> 预测类别: %s%n", 
                test[0], test[1], pred);
        }
    }
}

测试入口：Main.java

public class Main {
    public static void main(String[] args) {
        System.out.println("=== k-d树多维空间索引测试 ===\n");

        // 测试1：2D基础操作
        System.out.println("--- 测试1：2D最近邻 ---");
        List<KDPoint> points2d = Arrays.asList(
            new KDPoint(new double[]{2, 3}),
            new KDPoint(new double[]{5, 4}),
            new KDPoint(new double[]{9, 6}),
            new KDPoint(new double[]{4, 7}),
            new KDPoint(new double[]{8, 1}),
            new KDPoint(new double[]{7, 2})
        );

        KDTree tree2d = new KDTree(2);
        tree2d.build(points2d);
        System.out.println("构建2D树，点数: " + tree2d.size());

        KDPoint target = new KDPoint(new double[]{9, 2});
        KDPoint nearest = tree2d.nearest(target);
        System.out.println("目标点: " + target);
        System.out.println("最近点: " + nearest + " (距离: " + target.dist(nearest) + ")");

        // 测试2：3D K近邻
        System.out.println("\n--- 测试2：3D K近邻 ---");
        List<KDPoint> points3d = new ArrayList<>();
        java.util.Random rand = new java.util.Random(42);
        for (int i = 0; i < 20; i++) {
            points3d.add(new KDPoint(new double[]{
                rand.nextDouble() * 100,
                rand.nextDouble() * 100,
                rand.nextDouble() * 100
            }));
        }

        KDTree tree3d = new KDTree(3);
        tree3d.build(points3d);

        KDPoint target3d = new KDPoint(new double[]{50, 50, 50});
        List<KDPoint> kNearest = tree3d.kNearest(target3d, 5);
        System.out.println("目标点: " + target3d);
        System.out.println("5个最近邻:");
        for (int i = 0; i < kNearest.size(); i++) {
            KDPoint p = kNearest.get(i);
            System.out.printf("  #%d: %s (距离: %.2f)%n", 
                i + 1, p, target3d.dist(p));
        }

        // 测试3：范围查询
        System.out.println("\n--- 测试3：正交范围查询 ---");
        List<KDPoint> inRange = tree3d.rangeQuery(
            new double[]{30, 30, 30},
            new double[]{70, 70, 70}
        );
        System.out.println("范围 [30,70]x[30,70]x[30,70] 内点数: " + inRange.size());
        inRange.forEach(p -> System.out.println("  " + p));

        // 测试4：大规模性能
        System.out.println("\n--- 测试4：性能测试 (10万点3D) ---");
        List<KDPoint> bigData = new ArrayList<>();
        for (int i = 0; i < 100000; i++) {
            bigData.add(new KDPoint(new double[]{
                rand.nextDouble() * 1000,
                rand.nextDouble() * 1000,
                rand.nextDouble() * 1000
            }));
        }

        long start = System.currentTimeMillis();
        KDTree bigTree = new KDTree(3);
        bigTree.build(bigData);
        long buildTime = System.currentTimeMillis() - start;
        System.out.println("构建10万点3D树耗时: " + buildTime + "ms");

        start = System.currentTimeMillis();
        for (int i = 0; i < 1000; i++) {
            bigTree.nearest(new KDPoint(new double[]{
                rand.nextDouble() * 1000,
                rand.nextDouble() * 1000,
                rand.nextDouble() * 1000
            }));
        }
        long queryTime = System.currentTimeMillis() - start;
        System.out.println("1000次最近邻查询耗时: " + queryTime + "ms");

        // 测试5：KNN分类器
        System.out.println("\n--- 测试5：KNN分类器 ---");
        KNNClassifier.main(args);
    }
}

复杂度分析

操作 平均时间复杂度 最坏时间复杂度
构建（平衡） (O(n \log n)) (O(n \log n))
插入 (O(\log n)) (O(n))
最近邻搜索 (O(\log n)) (O(n))
K近邻搜索 (O(k \log n)) (O(n \log k))
范围查询 (O(m + \log n)) (O(n))

空间复杂度：(O(n))，每个点存储一次。

搜索剪枝原理

最近邻搜索的高效性来源于几何剪枝。假设当前最优距离为 (d)，目标点到当前节点划分超平面的距离为 (\Delta)：
– 若 (\Delta^2 \geq d^2)，则远侧子树中不可能存在更近的点，直接剪枝
– 否则，必须搜索远侧子树

这一剪枝策略使得平均搜索路径长度仅为 (O(\log n))，而非 (O(n))。

适用场景

  1. 机器学习：KNN分类/回归、特征向量检索
  2. 计算机视觉：SIFT/SURF特征点匹配
  3. 地理信息：多维地理坐标查询
  4. 推荐系统：用户/物品向量相似度搜索
  5. 物理模拟：粒子系统的近邻搜索

与四叉树的对比

特性 k-d树 四叉树
划分方式 按维度循环二分 固定四等分
维度扩展 任意k维 主要2D（3D为八叉树）
构建策略 自顶向下（静态） 自底向上（动态）
最近邻精度 精确 精确
最佳场景 高维KNN、静态数据 动态插入、范围查询

扩展方向

  1. 平衡k-d树：定期重建以维持平衡
  2. BBF（Best-Bin-First）：近似最近邻，优先队列控制搜索范围
  3. FLANN/Annoy：大规模高维近似最近邻库
  4. 球树（Ball Tree）：用超球面划分，在高维某些分布下更优

总结

k-d树是多维空间搜索的经典数据结构。通过Java实现，我们理解了按维度循环二分构建二叉树的优雅思想，以及基于超平面距离的几何剪枝策略。在机器学习KNN、特征匹配、空间数据库等需要高效多维查询的场景中，k-d树以其简洁的结构和优秀的平均性能，成为空间索引的核心工具之一。掌握k-d树，是深入理解空间算法和机器学习工程实践的重要一步。

Tagged with: