Algorithms-Part1最后一周的作业——KdTree

上两周的作业都没有写总结。今天一起完成吧。正好借此机会对这个课程的几个作业都回顾一下。
这次的作业陆陆续续写了一周，有两个原因吧。第一个就是Red-Black BST的应用比想象中的难一些，还有一个原因就是对递归的应用不熟。最初接触的时候以为很简单，像以前那样调用这周课学的数据结构。但想了想后就发现需要实现Red-Black BST，因为两个功能都要在树的搜索过程中完成。Insert函数的实现走了一些弯路，而且最后两天就是栽在这个函数上。因为在Insert一个node的时候需要根据父node设置自己的separator，递归写得不是那么好。但是为什么会导致最后奇怪的错误到现在也一直不清楚。只能下次避免写这类代码了。
Draw函数是optional的，本来没打算实现，因为想了几个方法都没法实现。还是在实现NearestPoint和Range的时候需要调试才实现的。切割Grid的方式就是把父Rect递归传入到子Grid。这不仅仅用来实现Draw，而且还是实现NearestPoint和Range的基本切分搜索思路。但是最后提交好作业后还有一些正确性是有问题的，但是发现不了。于是放弃了，知道思路就可以了。

import edu.princeton.cs.algs4.StdOut;
import edu.princeton.cs.algs4.Point2D;
import edu.princeton.cs.algs4.In;
import edu.princeton.cs.algs4.StdDraw;
import edu.princeton.cs.algs4.RectHV;
import edu.princeton.cs.algs4.Queue;

public class KdTree {

    private final static boolean VERTICAL = true;
    private final static boolean HORIZONTAL = false;
    private final static boolean RED = true;
    private final static boolean BLACK = false;

    private Node root;
    private int size;
    private double closestPresent;
    private Point2D closestPoint;

    public  KdTree() {
        root = null;
        size = 0;
    }

    private class Node {

        public Node left;
        public Node right;
        public Point2D key;
        public boolean separator;
        public boolean color;

        public Node(Point2D key, boolean color) {
            this.key = key;
            this.separator = separator;
        }
    }

    private Node rotateLeft(Node h) {
        assert h.right.color == RED;
        Node x = h.right;
        h.right = x.left;
        x.left = h;
        x.color = h.color;
        h.color = RED;
        return x;
    }

    private Node rotateRight(Node h) {
        assert h.left.color == RED;
        Node x = h.left;
        h.left = x.right;
        x.right = h;
        x.color = h.color;
        h.color = RED;
        return x;
    }

    private void flipColors(Node h) {
        assert h.left.color == RED;
        assert h.right.color == RED;
        assert h.color == BLACK;
        h.left.color = BLACK;
        h.right.color = BLACK;
        h.color = RED;
    }

    private boolean isRed(Node h) {
        if(h == null)
            return false;
        return h.color;
    }

    public boolean isEmpty() {
        return root == null;
    }

    public int size() {
        return this.size;
    }

    private int _count(Node n) {
        if(n == null)
            return 0;
        return _count(n.left)+_count(n.right)+1;
    }
    private int count(){
        return _count(root);
    }

    public void insert(Point2D p) {
        if(p == null)
            throw new IllegalArgumentException();
        if(root == null) {
            root = new Node(p, VERTICAL);
            size++;
            return;
        }
        if(contains(p))
            return;
        _insert(root, p);
        size++;
    }

    private Node _insert(Node h, Point2D p) {
        if(h == null)
            return new Node(p, RED);
        double cmp;
        if (h.separator == VERTICAL) {
            cmp = p.x() - h.key.x();
        } else {
            cmp = p.y() - h.key.y();
        }

        if(cmp <= 0) {
            h.left = _insert(h.left, p);
            h.left.separator = !h.separator;
        } else {
            h.right = _insert(h.right, p);
            h.right.separator = !h.separator;
        }

        if(isRed(h.right) && !isRed(h.left))
            h = rotateLeft(h);
        if(isRed(h.left) && isRed(h.left.left))
            h = rotateRight(h);
        if(isRed(h.right) && isRed(h.left))
            flipColors(h);
        return h;
    }

    public boolean contains(Point2D p) {
        return _contains(root, p);
    }

    private boolean _contains(Node n, Point2D p) {
        if(n == null)
            return false;
        double cmp;
        if(n.separator == VERTICAL) {
            cmp = p.x() - n.key.x();
        }else{
            cmp = p.y() - n.key.y();
        }
        if(n.key.equals(p))
            return true;
        if(cmp <= 0.0){
            return _contains(n.left, p);
        }else{
            return _contains(n.right, p);
        }
    }

    public void draw() {
        RectHV outer = new RectHV(0,0,1,1);
        _draw(root, outer);
    }

    private void _draw(Node n, RectHV outer) {
        if(n == null) return;
        if(n.separator == VERTICAL) {
            StdDraw.setPenColor(StdDraw.RED);
            StdDraw.setPenRadius(0.001);
            StdDraw.line(n.key.x(), outer.ymin(), n.key.x(), outer.ymax());
            StdDraw.setPenColor(StdDraw.BLACK);
            StdDraw.setPenRadius(0.01);
            StdDraw.point(n.key.x(), n.key.y());
            RectHV left = new RectHV(outer.xmin(),outer.ymin(),n.key.x(),outer.ymax());
            RectHV right = new RectHV(n.key.x(),outer.ymin(),outer.xmax(),outer.ymax());
            _draw(n.left, left);
            _draw(n.right, right);
        }else{
            StdDraw.setPenColor(StdDraw.BLUE);
            StdDraw.setPenRadius(0.001);
            StdDraw.line(outer.xmin(), n.key.y(), outer.xmax(), n.key.y());
            StdDraw.setPenColor(StdDraw.BLACK);
            StdDraw.setPenRadius(0.01);
            StdDraw.point(n.key.x(), n.key.y());
            RectHV down = new RectHV(outer.xmin(),outer.ymin(),outer.xmax(),n.key.y());
            RectHV up = new RectHV(outer.xmin(),n.key.y(),outer.xmax(),outer.ymax());
            _draw(n.left, down);
            _draw(n.right, up);
        }
    }

    public Iterable<Point2D> range(RectHV rect)  {
        if(rect == null)
            throw new IllegalArgumentException();
        Queue<Point2D> qu = new Queue<>();
        RectHV outer = new RectHV(0,0,1,1);
        _range(qu, root, rect, outer);
        return qu;
    }

    private void _range(Queue<Point2D> qu, Node n, RectHV rect, RectHV outer) {
        if(n == null) return;
        if(rect.contains(n.key))
            qu.enqueue(n.key);
        if(n.separator == VERTICAL) {
            RectHV left = new RectHV(outer.xmin(),outer.ymin(),n.key.x(),outer.ymax());
            RectHV right = new RectHV(n.key.x(),outer.ymin(),outer.xmax(),outer.ymax());
            if(rect.intersects(left))
                _range(qu, n.left, rect, outer);
            if(rect.intersects(right))
                _range(qu, n.right, rect, outer);
        }else{
            RectHV down = new RectHV(outer.xmin(),outer.ymin(),outer.xmax(),n.key.y());
            RectHV up = new RectHV(outer.xmin(),n.key.y(),outer.xmax(),outer.ymax());
            if(rect.intersects(down))
                _range(qu, n.left, rect, outer);
            if(rect.intersects(up))
                _range(qu, n.right, rect, outer);
        }
    }

    public Point2D nearest(Point2D p)   {
        if(p == null || root == null)
            throw new IllegalArgumentException();
        closestPresent = p.distanceTo(root.key);
        closestPoint = root.key;
        _nearest(root, p);
        return closestPoint;
    }

    private void _nearest(Node n, Point2D p) {
        if(n == null) return;
        double dis = p.distanceTo(n.key);
        if(dis < closestPresent) {
            closestPresent = dis;
            closestPoint = n.key;
        }
        if(n.separator == VERTICAL){
            if(p.x() <= n.key.x()){
                _nearest(n.left, p);
                if(closestPresent > p.x() - n.key.x()) {
                    _nearest(n.right, p);
                }
            } else {
                _nearest(n.right, p);
                if(closestPresent > n.key.x() - p.x()) {
                    _nearest(n.left, p);
                }
            }
        }else{
            if(p.y() <= n.key.y()) {
                _nearest(n.left, p);
                if(closestPresent > p.y() - n.key.y()) {
                    _nearest(n.right, p);
                }
            }
            else {
                _nearest(n.right, p);
                if(closestPresent > n.key.y() - p.y()) {
                    _nearest(n.left, p);
                }
            }
        }
    }

    public static void main(String[] args) {
        KdTree sets = new KdTree();
        In in = new In(args[0]);
        while(!in.isEmpty()){
            double x = in.readDouble();
            double y = in.readDouble();
            sets.insert(new Point2D(x,y));
        }
        in.close();
        StdOut.println("Size: " + sets.size());
    }
}