每日算法 — 使用java实现KD树：k维空间划分与最近邻搜索

KD树（K-Dimensional Tree）是一种高效的空间划分数据结构，广泛应用于最近邻搜索、范围查询与多维数据索引。本文用Java实现KD树，讲解k维空间递归划分、最近邻搜索与范围查询算法。

核心数据结构

import java.util.*;

public class KDTree {
    // 二维点（可扩展为任意维度）
    static class Point {
        double[] coords;
        Point(double... coords) { this.coords = coords; }
        int dim() { return coords.length; }
        double get(int i) { return coords[i]; }

        double squaredDistance(Point other) {
            double sum = 0;
            for (int i = 0; i < dim(); i++) {
                double d = coords[i] - other.coords[i];
                sum += d * d;
            }
            return sum;
        }

        @Override
        public String toString() {
            return Arrays.toString(coords);
        }
    }

    // KD树节点
    static class Node {
        Point point;
        Node left, right;
        int axis; // 当前节点的划分维度

        Node(Point point, int axis) {
            this.point = point;
            this.axis = axis;
        }
    }

    Node root;
    int dimension;

    KDTree(int dimension) {
        this.dimension = dimension;
    }

递归构建：中位数划分

    // 构建KD树：按中位数递归划分，保证树高O(log n)
    void build(List<Point> points) {
        root = buildRecursive(points, 0);
    }

    Node buildRecursive(List<Point> points, int depth) {
        if (points.isEmpty()) return null;
        int axis = depth % dimension;
        // 按当前维度排序并取中位数
        points.sort(Comparator.comparingDouble(p -> p.get(axis)));
        int median = points.size() / 2;
        Node node = new Node(points.get(median), axis);
        node.left = buildRecursive(points.subList(0, median), depth + 1);
        node.right = buildRecursive(points.subList(median + 1, points.size()), depth + 1);
        return node;
    }

k近邻与范围查询

    // k近邻搜索：使用优先队列（大顶堆）维护最近的k个点
    List<Point> kNearest(Point target, int k) {
        PriorityQueue<Result> pq = new PriorityQueue<>(
            Comparator.comparingDouble((Result r) -> r.dist).reversed());
        kNearest(root, target, k, pq);
        List<Point> result = new ArrayList<>();
        for (Result r : pq) result.add(r.point);
        return result;
    }

    void kNearest(Node node, Point target, int k, PriorityQueue<Result> pq) {
        if (node == null) return;
        double dist = node.point.squaredDistance(target);
        if (pq.size() < k) {
            pq.offer(new Result(node.point, dist));
        } else if (dist < pq.peek().dist) {
            pq.poll();
            pq.offer(new Result(node.point, dist));
        }

        int axis = node.axis;
        double diff = target.get(axis) - node.point.get(axis);
        Node near = diff < 0 ? node.left : node.right;
        Node far = diff < 0 ? node.right : node.left;

        kNearest(near, target, k, pq);
        if (diff * diff < (pq.size() < k ? Double.MAX_VALUE : pq.peek().dist)) {
            kNearest(far, target, k, pq);
        }
    }

    // 范围查询：返回矩形区域内的所有点
    List<Point> rangeQuery(Point min, Point max) {
        List<Point> result = new ArrayList<>();
        rangeQuery(root, min, max, result);
        return result;
    }

    void rangeQuery(Node node, Point min, Point max, List<Point> result) {
        if (node == null) return;
        if (inRange(node.point, min, max)) result.add(node.point);
        int axis = node.axis;
        if (min.get(axis) <= node.point.get(axis))
            rangeQuery(node.left, min, max, result);
        if (max.get(axis) >= node.point.get(axis))
            rangeQuery(node.right, min, max, result);
    }

    boolean inRange(Point p, Point min, Point max) {
        for (int i = 0; i < dimension; i++) {
            if (p.get(i) < min.get(i) || p.get(i) > max.get(i)) return false;
        }
        return true;
    }

主程序与性能测试

    public static void main(String[] args) {
        int n = 100000;
        Random rand = new Random(42);
        List<Point> points = new ArrayList<>();
        for (int i = 0; i < n; i++) {
            points.add(new Point(rand.nextDouble() * 1000, 
                                 rand.nextDouble() * 1000));
        }

        KDTree tree = new KDTree(2);
        long t1 = System.currentTimeMillis();
        tree.build(new ArrayList<>(points));
        long t2 = System.currentTimeMillis();
        System.out.println("构建KD树耗时: " + (t2 - t1) + "ms");

        // 最近邻搜索测试
        Point target = new Point(500, 500);
        t1 = System.currentTimeMillis();
        Point nearest = tree.nearestNeighbor(target);
        t2 = System.currentTimeMillis();
        System.out.println("最近邻: " + nearest + ", 距离: " + 
            Math.sqrt(target.squaredDistance(nearest)));
        System.out.println("最近邻搜索耗时: " + (t2 - t1) + "ms");

        // 暴力对比
        t1 = System.currentTimeMillis();
        Point brute = points.get(0);
        double minDist = target.squaredDistance(brute);
        for (Point p : points) {
            double d = target.squaredDistance(p);
            if (d < minDist) { minDist = d; brute = p; }
        }
        t2 = System.currentTimeMillis();
        System.out.println("暴力搜索耗时: " + (t2 - t1) + "ms");

        // k近邻测试
        List<Point> knn = tree.kNearest(target, 5);
        System.out.println("5近邻: " + knn);
    }
}

算法分析

构建复杂度：每层排序 O(n log n)，递归深度 O(log n)，总构建时间 O(n log^2 n)，若使用线性时间中位数算法可优化至 O(n log n)。

最近邻搜索：平均 O(log n)，最坏 O(n)。剪枝条件 diff^2 < bestDist 是关键：若目标点到划分超平面的距离大于当前最优距离，则无需搜索另一侧子树。

空间复杂度：存储所有点 O(n)。

Tagged with: Java KD树数据结构最近邻搜索空间划分