如何在 Java 中基于 KD 树高效实现 M 近邻搜索（k-NN 扩展版）

花韻仙語

发布时间：2026-01-19 09:55:11

152人浏览过

来源于php中文网

原创

如何在 Java 中基于 KD 树高效实现 M 近邻搜索（k-NN 扩展版）

本文详解如何在不依赖第三方库的前提下，基于自定义 kd 树结构，用 java 实现 `float[][] findmnearest(float[] point, int m)` 方法，支持返回距离查询点最近的 m 个样本坐标，涵盖剪枝策略、最大堆优化与递归回溯逻辑。

在单近邻（1-NN）搜索中，我们维护一个全局最优节点 best 和最小距离 bestDistance，通过轴对齐分割与超矩形剪枝高效遍历。但扩展到 M 近邻（M-NN） 时，核心挑战在于：
✅ 不再只需跟踪“当前最近”，而需动态维护 候选集 Top-M；
✅ 剪枝条件必须升级——不能仅比较单点距离，而需判断「当前子树是否可能包含比当前第 M 近点更近的点」；
✅ 需避免重复访问或遗漏，尤其在回溯时需重新评估另一侧子树。

✅ 解决方案：最大堆 + 递归回溯（推荐）

使用 PriorityQueue 构建固定容量的最大堆（按欧氏距离平方排序），始终保留距离最小的 m 个点。堆顶即当前第 m 近的距离上限 maxHeapTopDistSq，用于关键剪枝：

import java.util.*;

public class KDTree {
    private static class KDNode {
        final float[] coords;
        KDNode left, right;
        int axis; // splitting axis (0, 1, ..., k-1)

        KDNode(float[] coords, int axis) {
            this.coords = coords.clone();
            this.axis = axis;
        }

        float distanceSq(KDNode other) {
            float sum = 0f;
            for (int i = 0; i < coords.length; i++) {
                float d = coords[i] - other.coords[i];
                sum += d * d;
            }
            return sum;
        }

        float getCoordinate(int dim) { return coords[dim]; }
    }

    private KDNode root;
    private final int k; // dimensionality

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

    // Main M-NN method
    public float[][] findMNearest(float[] point, int m) {
        if (point == null || m <= 0 || root == null) 
            return new float[0][0];

        // Max-heap: store [distance^2, coordinates] → sort by distance^2 descending
        PriorityQueue maxHeap = new PriorityQueue<>((a, b) -> 
            Float.compare(b[0], a[0]) // descending order
        );

        // Recursive search with pruning
        searchMNN(root, new KDNode(point, 0), 0, maxHeap, m);

        // Extract top m points (heap may contain < m if tree size < m)
        float[][] result = new float[maxHeap.size()][k];
        int i = 0;
        while (!maxHeap.isEmpty()) {
            float[] entry = maxHeap.poll();
            System.arraycopy(entry, 1, result[i++], 0, k); // skip dist at index 0
        }
        return result;
    }

    private void searchMNN(KDNode node, KDNode target, int depth, 
                          PriorityQueue heap, int m) {
        if (node == null) return;

        int axis = depth % k;
        float distSq = node.distanceSq(target);
        float[] coords = node.coords;

        // Insert current node if heap not full, or replace worst if closer
        if (heap.size() < m) {
            float[] entry = new float[k + 1];
            entry[0] = distSq;
            System.arraycopy(coords, 0, entry, 1, k);
            heap.offer(entry);
        } else if (distSq < heap.peek()[0]) {
            heap.poll(); // remove worst
            float[] entry = new float[k + 1];
            entry[0] = distSq;
            System.arraycopy(coords, 0, entry, 1, k);
            heap.offer(entry);
        }

        // Determine which child is closer & visit first (better pruning chance)
        boolean goLeftFirst = (target.getCoordinate(axis) < node.getCoordinate(axis));
        KDNode nearChild = goLeftFirst ? node.left : node.right;
        KDNode farChild  = goLeftFirst ? node.right : node.left;

        // Visit near subtree first
        searchMNN(nearChild, target, depth + 1, heap, m);

        // Pruning: check if far subtree can contain better candidates
        float diff = target.getCoordinate(axis) - node.getCoordinate(axis);
        float diffSq = diff * diff;

        // If heap is not full, we MUST check far side (no pruning)
        // If heap is full, only explore far side if diffSq < heap's max distance^2
        if (heap.size() == m && diffSq < heap.peek()[0]) {
            searchMNN(farChild, target, depth + 1, heap, m);
        }
    }
}