kdtree.java

2004年robotcup世界冠军源代码

package MRL.Utilities;

import java.awt.*;
import java.awt.geom.*;
import java.util.*;
import yab.agent.object.*;


public class KdTree {
    private KdTree cl;
    private KdTree cr;


    private MotionlessObject obj;
    private Rectangle2D rgn;

    private class Cell {
        final MotionlessObject obj;
        Cell next;
        Cell(MotionlessObject o, Cell n) {
            obj = o;
            next = n;
        }

        public String toString() {
            return "(" + obj.x() + "," + obj.y() + ")";
        }
    }


    private static final Comparator X_CMP = new Comparator() {
        public int compare(Object o1, Object o2) {
            int x1 = ((Cell) o1).obj.x();
            int x2 = ((Cell) o2).obj.x();
            int comparison = x1 - x2;
            if (comparison != 0) {
                return comparison;
            }
            int y1 = ((Cell) o1).obj.y();
            int y2 = ((Cell) o2).obj.y();
            return y1 - y2;
        }

        public boolean equals(Object obj) {
            return super.equals(obj);
        }
    };


    private static final Comparator Y_CMP = new Comparator() {
        public int compare(Object o1, Object o2) {
            int y1 = ((Cell) o1).obj.y();
            int y2 = ((Cell) o2).obj.y();
            int comparison = y1 - y2;
            if (comparison != 0) {
                return comparison;
            }
            int x1 = ((Cell) o1).obj.x();
            int x2 = ((Cell) o2).obj.x();
            return x1 - x2;
        }

        public boolean equals(Object obj) {
            return super.equals(obj);
        }
    };


    private class AList {
        final int size;
        final Cell[] array;
        final Cell head;


        AList(Set set, Comparator cmp) {
            size = set.size();
            array = new Cell[size];
            Iterator it = set.iterator();
            for (int i = 0; i < array.length; i++) {
                array[i] = new Cell((MotionlessObject) it.next(), null);
            }
            Arrays.sort(array, cmp);
            head = array[0];
            for (int i = 0; i < array.length - 1; i++) {
                array[i].next = array[i + 1];
            }
        }

        AList(Cell[] a, Cell h, int s) {
            array = a;
            head = h;
            size = s;
        }


        Object[] divide() {
            Object[] result = new Object[3];
            Cell tmp = head;
            int mid = (size + 1) / 2;
            for (int i = 0; i < mid - 1; i++) {
                tmp = tmp.next;
            }
            result[0] = tmp;
            result[2] = new AList(array, tmp.next, size - mid);
            tmp.next = null;
            result[1] = new AList(array, head, mid);
            return result;
        }


        Object[] divide(Cell mid, Comparator cmp) {
            Cell hl = null, hr = null, tl = null, tr = null;
            for (Cell p = head; p != null; p = p.next) {
                if (cmp.compare(p, mid) <= 0) {
                    if (hl == null) {
                        hl = p;
                    } else {
                        tl.next = p;
                    }
                    tl = p;
                } else {
                    if (hr == null) {
                        hr = p;
                    } else {
                        tr.next = p;
                    }
                    tr = p;
                }
            }
            tl.next = null;
            tr.next = null;
            int s = (size + 1) / 2;
            return new Object[] {
                    new AList(array, hl, s), new AList(array, hr, size - s)};
        }

        public String toString() {
            String result = "(size: " + size;
            for (Cell p = head; p != null; p = p.next) {
                result += "  " + p;
            }
            return result + ")";
        }
    }


    public KdTree(Set motionlessObjectSet) {
        this();

        AList alx = new AList(motionlessObjectSet, X_CMP);
        AList aly = new AList(motionlessObjectSet, Y_CMP);
        KdTree tmp = build(alx, aly, 0, MAX_FIELD);
        cl = tmp.cl;
        cr = tmp.cr;
        obj = tmp.obj;
        rgn = tmp.rgn;
    }

    private KdTree() {
        this(null, null, null, null);
    }

    private KdTree(KdTree l, KdTree r, MotionlessObject o, Rectangle2D region) {
        cl = l;
        cr = r;
        obj = o;
        rgn = region;
    }

    private KdTree build(AList alx, AList aly, int depth, Rectangle2D rgn) {
        if (alx.size == 1) { // alx.size == aly.size
            return new KdTree(null, null, alx.head.obj, null);
        }

        KdTree result = new KdTree();
        if ((depth & 0x1) == 0) {

            Object[] tmp;
            tmp = alx.divide();
            result.obj = ((Cell) tmp[0]).obj;
            AList xl = (AList) tmp[1];
            AList xr = (AList) tmp[2];
            tmp = aly.divide((Cell) tmp[0], X_CMP);
            AList yl = (AList) tmp[0];
            AList yr = (AList) tmp[1];
            Rectangle2D rgnl = regionl(rgn, depth, result.obj);
            Rectangle2D rgnr = regionr(rgn, depth, result.obj);
            result.cl = build(xl, yl, depth + 1, rgnl);
            result.cr = build(xr, yr, depth + 1, rgnr);
            result.rgn = rgn;
        } else {

            Object[] tmp;
            tmp = aly.divide();
            result.obj = ((Cell) tmp[0]).obj;
            AList yl = (AList) tmp[1];
            AList yr = (AList) tmp[2];
            tmp = alx.divide((Cell) tmp[0], Y_CMP);
            AList xl = (AList) tmp[0];
            AList xr = (AList) tmp[1];
            Rectangle2D rgnl = regionl(rgn, depth, result.obj);
            Rectangle2D rgnr = regionr(rgn, depth, result.obj);
            result.cl = build(xl, yl, depth + 1, rgnl);
            result.cr = build(xr, yr, depth + 1, rgnr);
            result.rgn = rgn;
        }
        return result;
    }

    private static final Rectangle2D MAX_FIELD
            = new Rectangle2D.Float(0, 0, Integer.MAX_VALUE, Integer.MAX_VALUE);
    private Rectangle2D m_rectangle = new Rectangle2D.Float();
    private Rectangle2D regionl(Rectangle2D rgn, int depth,
                                MotionlessObject obj) {
        if ((depth & 0x1) == 0) {
            m_rectangle.setRect(0, 0, obj.x(), Integer.MAX_VALUE);
        } else {
            m_rectangle.setRect(0, 0, Integer.MAX_VALUE, obj.y());
        }
        return m_rectangle.createIntersection(rgn);
    }

    private Rectangle2D regionr(Rectangle2D rgn, int depth,
                                MotionlessObject obj) {
        if ((depth & 0x1) == 0) {

            m_rectangle.setRect(obj.x() + 1, 0, Integer.MAX_VALUE,
                                Integer.MAX_VALUE);
        } else {

            m_rectangle.setRect(0, obj.y() + 1, Integer.MAX_VALUE,
                                Integer.MAX_VALUE);
        }
        return m_rectangle.createIntersection(rgn);
    }

    private final Ellipse2D m_ellipse = new Ellipse2D.Float();
    private Ellipse2D ellipse(RealObject obj, int r) {
        return ellipse(obj.x(), obj.y(), r);
    }

    private Ellipse2D ellipse(int x, int y, int r) {
        m_ellipse.setFrame(x - r, y - r, r * 2, r * 2);
        return m_ellipse;
    }

    private final Point2D m_point = new Point2D.Float();
    private Point2D point(RealObject obj) {
        return point(obj.x(), obj.y());
    }

    private Point2D point(int x, int y) {
        m_point.setLocation(x, y);
        return m_point;
    }

    private boolean isLeaf() {
        return (cl == null && cr == null);
    }


    public ArrayList search(RealObject obj, int r) {
        return search(ellipse(obj, r));
    }


    public ArrayList search(Shape r) {
        ArrayList result = new ArrayList();
        search(this, new Area(r), result);
        return result;
    }

    private void search(KdTree v, Area r, ArrayList set) {
        if (v.isLeaf()) {
            if (r.contains(point(v.obj))) {
                set.add(v.obj);
            }
            return;
        }

        if (r.contains(v.rgn)) {
            reportSubtree(v.cl, set);
        } else {
            if (r.intersects(v.rgn)) {
                search(v.cl, r, set);
            }
        }

        if (r.contains(v.rgn)) {
            reportSubtree(v.cr, set);
        } else {
            if (r.intersects(v.rgn)) {
                search(v.cr, r, set);
            }
        }
    }

    private void reportSubtree(KdTree v, ArrayList set) {
        if (v.isLeaf()) {
            set.add(v.obj);
            return;
        }
        reportSubtree(v.cl, set);
        reportSubtree(v.cr, set);
    }
}