Data Structure - K-D Tree (Java)

版权声明：欢迎转载并请注明出处，谢谢~~ https://blog.csdn.net/chimomo/article/details/85317986

package chimomo.learning.java.datastructure;

/**
 * Quick illustration of a two-dimensional tree.
 *
 * @author Created by Chimomo
 */
public class KdTree<T extends Comparable<? super T>> {
    private KdNode<T> root;

    public KdTree() {
        root = null;
    }

    public static void main(String[] args) {
        // Construct K-D tree.
        KdTree<Integer> t = new KdTree<>();

        // Insert.
        for (int i = 300; i < 370; i++) {
            Integer[] arr = new Integer[2];
            arr[0] = i;
            arr[1] = 2500 - i;
            t.insert(arr);
        }

        // Print range.
        Integer[] low = {70, 2186};
        Integer[] high = {1200, 2200};
        t.printRange(low, high);
    }

    public void insert(T[] x) {
        root = insert(x, root, 0);
    }

    private KdNode<T> insert(T[] x, KdNode<T> t, int level) {
        if (t == null) {
            t = new KdNode<>(x);
        } else if (x[level].compareTo(t.data[level]) < 0) {
            t.left = insert(x, t.left, 1 - level);
        } else {
            t.right = insert(x, t.right, 1 - level);
        }
        return t;
    }

    /**
     * Print items satisfying low[ 0 ] <= x[ 0 ] <= high[ 0 ] and low[ 1 ] <= x[ 1 ] <= high[ 1 ].
     */
    public void printRange(T[] low, T[] high) {
        printRange(low, high, root, 0);
    }

    private void printRange(T[] low, T[] high, KdNode<T> t, int level) {
        if (t != null) {
            if (low[0].compareTo(t.data[0]) <= 0
                    && low[1].compareTo(t.data[1]) <= 0
                    && high[0].compareTo(t.data[0]) >= 0
                    && high[1].compareTo(t.data[1]) >= 0) {
                System.out.println("(" + t.data[0] + "," + t.data[1] + ")");
            }
            if (low[level].compareTo(t.data[level]) <= 0) {
                printRange(low, high, t.left, 1 - level);
            }
            if (high[level].compareTo(t.data[level]) >= 0) {
                printRange(low, high, t.right, 1 - level);
            }
        }
    }

    private static class KdNode<T> {
        T[] data;
        KdNode<T> left;
        KdNode<T> right;

        KdNode(T item[]) {
            data = (T[]) new Comparable[2];
            data[0] = item[0];
            data[1] = item[1];
            left = right = null;
        }
    }
}

/*
Output:
(300,2200)
(301,2199)
(302,2198)
(303,2197)
(304,2196)
(305,2195)
(306,2194)
(307,2193)
(308,2192)
(309,2191)
(310,2190)
(311,2189)
(312,2188)
(313,2187)
(314,2186)

*/