QuadTree.java

package org.djutils.quadtree;

import java.io.Serializable;
import java.util.Collection;
import java.util.Iterator;
import java.util.LinkedHashSet;
import java.util.Set;

import org.djutils.exceptions.Throw;

/**
 * Quad tree for 2D objects. For now, this implementation needs an ultimate outer bounding box. No part of any 2D string object
 * may exceed that bounding box. A link to each stored 2D object will be stored in each sub-box that it intersects.
 * <p>
 * @author <a href="http://www.tudelft.nl/pknoppers">Peter Knoppers</a>
 * @param <T> Type of object stored in this quad tree
 */
public class QuadTree<T extends Envelope> implements Collection<T>, Serializable
{
    /** ... */
    private static final long serialVersionUID = 20200904L;

    /** Maximum number of payload objects in one cell. */
    private final int maximumLoad;

    /** Minimum width and height of a SubTree bounding box. */
    private final double minimumSize;

    /** The actual top level quad tree. */
    private final SubTree<T> tree;

    /** Count the number of sub trees created. */
    private int totalSubTrees = 0;

    /**
     * Create a new QuadTree object (or a sub-tree).
     * @param maximumLoad int; number of elements at any level that warrants investigating if the tree can be re-balanced
     * @param minimumSize double; minimum width or height of a sub tree Rectangle (smaller sub tree are never created)
     * @param left double; the lowest X-coordinate that is allowed (inclusive)
     * @param bottom double; the lowest Y-coordinate that is allowed (inclusive)
     * @param right double; the highest X-coordinate that is allowed (exclusive)
     * @param top double; the highest Y-coordinate that is allowed (exclusive)
     */
    public QuadTree(final int maximumLoad, final double minimumSize, final double left, final double bottom, final double right,
            final double top)
    {
        Throw.when(left >= right, IllegalArgumentException.class, "left (%f) must be less than right (%f)", left, right);
        Throw.when(bottom >= top, IllegalArgumentException.class, "bottom (%f) must be less than top (%f)", bottom, top);
        this.maximumLoad = maximumLoad;
        this.minimumSize = minimumSize;
        this.tree = new SubTree<T>(this, new Rectangle(left, bottom, right, top));
    }

    /**
     * Return the number of objects at which it is time to try to re-balance.
     * @return int; the number of objects at which it is time to try to re-balance
     */
    public int getMaxLoad()
    {
        return this.maximumLoad;
    }

    /**
     * Return the minimum sub-tree rectangle size.
     * @return double; the minimum sub-tree rectangle size
     */
    public double getMinimumSize()
    {
        return this.minimumSize;
    }

    /** {@inheritDoc} */
    @Override
    public int size()
    {
        return this.tree.size();
    }

    /** {@inheritDoc} */
    @Override
    public boolean isEmpty()
    {
        return this.tree.size() == 0;
    }

    /** {@inheritDoc} */
    @Override
    public boolean contains(final Object o)
    {
        if (!(o instanceof Envelope))
        {
            return false;
        }
        @SuppressWarnings("unchecked")
        T t = (T) o;
        return this.tree.recursiveContains(new RectangleAndPayload<T>(t.getBoundingRectangle(), t));
    }

    /** {@inheritDoc} */
    @Override
    public Iterator<T> iterator()
    {
        return iterator(this.tree.getBoundingBox());
    }

    /**
     * Find all elements intersecting a given bounding box. This iterator cannot be used to remove elements, but the remove
     * method can be safely called while the iterator is active.
     * @param searchArea Rectangle; the bounding box
     * @return Iterator&lt;T&gt;; iterator that returns all elements that intersect the given bounding box
     */
    public Iterator<T> iterator(final Rectangle searchArea)
    {
        return collect(searchArea).iterator();
    }

    /** {@inheritDoc} */
    @Override
    public Object[] toArray()
    {
        return collect(this.tree.getBoundingBox()).toArray();
    }

    /** {@inheritDoc} */
    @SuppressWarnings("hiding")
    @Override
    public <T> T[] toArray(final T[] a)
    {
        return collect(this.tree.getBoundingBox()).toArray(a);
    }

    /**
     * Construct a set containing all payload elements within a specified area.
     * @param searchArea Rectangle; the search area
     * @return Set&lt;T&gt;; the set containing all payload elements whose bounding areas intersect the specified area
     */
    private Set<T> collect(final Rectangle searchArea)
    {
        Iterator<RectangleAndPayload<T>> iterator = this.tree.recursiveCollect(searchArea).iterator();
        Set<T> result = new LinkedHashSet<>();
        while (iterator.hasNext())
        {
            result.add(iterator.next().getPayload());
        }
        return result;
    }

    /** {@inheritDoc} */
    @Override
    public boolean add(final T e)
    {
        return this.tree.add(new RectangleAndPayload<T>(e.getBoundingRectangle(), e));
    }

    /** {@inheritDoc} */
    @Override
    public boolean remove(final Object o)
    {
        if (!(o instanceof Envelope))
        {
            return false;
        }
        @SuppressWarnings("unchecked")
        T t = (T) o;
        return this.tree.remove(new RectangleAndPayload<T>(t.getBoundingRectangle(), t));
    }

    /** {@inheritDoc} */
    @Override
    public boolean containsAll(final Collection<?> c)
    {
        return collect(this.tree.getBoundingBox()).containsAll(c);
    }

    /** {@inheritDoc} */
    @Override
    public boolean addAll(final Collection<? extends T> c)
    {
        boolean result = false;
        for (T t : c)
        {
            result |= add(t);
        }
        return result;
    }

    /** {@inheritDoc} */
    @Override
    public boolean removeAll(final Collection<?> c)
    {
        boolean result = false;
        for (Object o : c)
        {
            result |= remove(o);
        }
        return result;
    }

    /** {@inheritDoc} */
    @Override
    public boolean retainAll(final Collection<?> c)
    {
        throw new RuntimeException("Not (yet) implemented");
    }

    /** {@inheritDoc} */
    @Override
    public void clear()
    {
        this.tree.clear();
    }

    /**
     * Increment the number of sub trees created.
     */
    void incrementSubTreeCount()
    {
        this.totalSubTrees++;
    }

    /**
     * Return the total number of sub trees.
     * @return int; the total number of sub trees
     */
    public int getSubTreeCount()
    {
        return this.totalSubTrees;
    }

    /** {@inheritDoc} */
    @Override
    public String toString()
    {
        return "QuadTree [maximumLoad=" + this.maximumLoad + ", minimumSize=" + this.minimumSize + ", tree=" + this.tree + "]";
    }

    /**
     * Make a textual description of this quad tree drilling down to the prescribed depth.
     * @param expandDepth int; maximum depth to descend
     * @return String; textual description of this quad tree
     */
    public String toString(final int expandDepth)
    {
        return "QuadTree [maximumLoad=" + this.maximumLoad + ", minimumSize=" + this.minimumSize + ", tree="
                + this.tree.toString(expandDepth) + "]";
    }

    /**
     * Dump a quad tree.
     * @param indent String; prefix for each output line
     * @return String; textual description of this quad tree.
     */
    public String dump(final String indent)
    {
        return this.tree.dump(indent);
    }

    /**
     * Sub tree of a quad tree.
     * @param <T> Type of object stored in this quad tree
     */
    @SuppressWarnings("hiding")
    class SubTree<T extends Envelope> implements Serializable
    {
        /** ... */
        private static final long serialVersionUID = 20200904L;

        /** Root of the quad tree. */
        private final QuadTree<T> root;

        /** Bounding box of this quad tree. */
        private final Rectangle boundingBox;

        /** Current number of objects in this quad tree. Includes all children, counting each object exactly once. */
        private int size = 0;

        /** If the four children have been allocated, this array will be non-null and contain the four children. */
        private SubTree<T>[] children = null;

        /** Elements stored at this node. */
        private Set<RectangleAndPayload<T>> elements = null;

        /**
         * Construct a new sub tree.
         * @param root QuadTree&lt;T&gt;; the root
         * @param boundingBox Rectangle; the bounding box of the new sub tree
         */
        SubTree(final QuadTree<T> root, final Rectangle boundingBox)
        {
            this.root = root;
            this.boundingBox = boundingBox;
            root.incrementSubTreeCount();
        }

        /**
         * Retrieve the bounding box of this sub tree.
         * @return Rectangle; the bounding box of this sub tree
         */
        public final Rectangle getBoundingBox()
        {
            return this.boundingBox;
        }

        /**
         * Return the number of objects stored in and under this SubTree.
         * @return int; the number of objects stored in and under this SubTree
         */
        public int size()
        {
            return this.size;
        }

        /**
         * Add a RectangleAndPayload to this SubTree.
         * @param e RectangleAndPayload&lt;T&gt;; the object to add
         * @return boolean; true if this SubTree was changed (object was added); false if this SubTree did not change
         */
        public boolean add(final RectangleAndPayload<T> e)
        {
            if (contains(e))
            {
                return false;
            }
            if (this.elements == null)
            {
                this.elements = new LinkedHashSet<>();
            }
            this.elements.add(e);
            this.size++;
            reBalance();
            return true;
        }

        /**
         * Remove a RectangleAndPayload from this SubTree.
         * @param o RectangleAndPayload&lt;T&gt;; the object to remove
         * @return boolean; true if this SubTree was changed (object was removed); false if this SubTree did not change
         */
        public boolean remove(final RectangleAndPayload<T> o)
        {
            if (this.elements.remove(o))
            {
                this.size--;
                return true; // The object cannot be also present in any of the sub trees
            }
            // Try all of the sub trees
            boolean result = false;
            if (this.children != null)
            {
                Rectangle rectangle = o.getRectangle();
                for (SubTree<T> child : this.children)
                {
                    if (!child.boundingBox.intersects(rectangle))
                    {
                        continue; // This is the time saver
                    }
                    if (child.remove(o))
                    {
                        result = true;
                    }
                }
            }
            if (result)
            {
                this.size--;
            }
            return result;
        }

        /**
         * Delete all objects stored in this SubTree.
         */
        public void clear()
        {
            this.elements.clear();
            this.children = null;
            this.size = 0;
        }

        /**
         * Determine if this SubTree contains a specific object.
         * @param o RectangleAndPayload&lt;T&gt;; the object to search
         * @return boolean; true if this SubTree contains the object
         */
        public boolean contains(final RectangleAndPayload<T> o)
        {
            if (this.elements == null)
            {
                return false;
            }
            return recursiveContains(o);
        }

        /**
         * Recursively search for a particular object.
         * @param o RectangleAndPayload&lt;T&gt;; the object to search for
         * @return boolean; true if this quad tree contains the object; false if this quad tree does not contain the object
         */
        boolean recursiveContains(final RectangleAndPayload<T> o)
        {
            if ((!this.boundingBox.intersects(o.getRectangle())) || (this.elements == null))
            {
                return false; // This is the time saver
            }
            for (RectangleAndPayload<T> element : this.elements)
            {
                if (element.equals(o))
                {
                    return true;
                }
            }
            if (this.children == null)
            {
                return false;
            }
            for (SubTree<T> child : this.children)
            {
                if (child.recursiveContains(o))
                {
                    return true;
                }
            }
            return false;
        }

        /**
         * Recursively collect all elements that intersect the given rectangle.
         * @param rectangle Rectangle; the rectangle
         * @return Set&lt;RectangleAndPayload&lt;T&gt;&gt;; all stored elements that intersect the given rectangle
         */
        public Set<RectangleAndPayload<T>> recursiveCollect(final Rectangle rectangle)
        {
            Set<RectangleAndPayload<T>> result = new LinkedHashSet<>();
            if (!this.boundingBox.intersects(rectangle))
            {
                return result; // This is the time saver
            }
            if (this.elements != null)
            {
                for (RectangleAndPayload<T> element : this.elements)
                {
                    if (element.getRectangle().intersects(rectangle))
                    {
                        result.add(element);
                    }
                }
            }
            if (this.children != null)
            {
                for (SubTree<T> child : this.children)
                {
                    result.addAll(child.recursiveCollect(rectangle));
                }
            }
            return result;
        }

        /**
         * Optimize the distribution of elements at this node and at sub-nodes.
         */
        @SuppressWarnings("unchecked")
        private void reBalance()
        {
            if (this.elements.size() < this.root.getMaxLoad() || this.boundingBox.getWidth() < this.root.getMinimumSize()
                    || this.boundingBox.getHeight() < this.root.getMinimumSize())
            {
                return;
            }
            // Count the number of elements that could be moved down to sub-trees
            double cX = (this.boundingBox.getLeft() + this.boundingBox.getRight()) / 2;
            double cY = (this.boundingBox.getBottom() + this.boundingBox.getTop()) / 2;
            int canMove = 0;
            /*-
            for (T e : this.elements)
            {
                // This criterion is not good
                if (!e.getBoundingRectangle().contains(this.boundingBox))
                {
                    canMove++;
                }
            }
            */
            canMove = this.elements.size();
            if (canMove == 0 || canMove < this.root.getMaxLoad() /* / 2 */ && this.children == null)
            {
                // System.out.println("reBalance: not moving " + canMove + " of " + this.elements.size());
                return;
            }
            // System.out.println("At start of reBalance of " + this.toString(1));
            if (this.children == null)
            {
                this.children = new SubTree[] {
                        new SubTree<T>(this.root,
                                new Rectangle(this.boundingBox.getLeft(), this.boundingBox.getBottom(), cX, cY)),
                        new SubTree<T>(this.root,
                                new Rectangle(cX, this.boundingBox.getBottom(), this.boundingBox.getRight(), cY)),
                        new SubTree<T>(this.root, new Rectangle(this.boundingBox.getLeft(), cY, cX, this.boundingBox.getTop())),
                        new SubTree<T>(this.root,
                                new Rectangle(cX, cY, this.boundingBox.getRight(), this.boundingBox.getTop()))};
            }
            Iterator<RectangleAndPayload<T>> iterator = this.elements.iterator();
            while (iterator.hasNext())
            {
                RectangleAndPayload<T> e = iterator.next();
                if (e.getRectangle().contains(this.boundingBox))
                {
                    continue;
                }
                boolean added = false;
                for (SubTree<T> child : this.children)
                {
                    if (e.getRectangle().intersects(child.boundingBox))
                    {
                        added |= child.add(e);
                    }
                }
                if (added)
                {
                    iterator.remove();
                }
                else
                {
                    System.out.println("ERROR: Could not add " + e + " to any of the children");
                }
            }
            // System.out.println("At end of reBalanceof " + this.toString(1));
        }

        /** {@inheritDoc} */
        @Override
        public String toString()
        {
            return "SubTree [boundingBox=" + this.boundingBox + ", size=" + this.size + ", children=" + this.children
                    + (this.elements == null ? "elements=null" : ", elements.size=" + this.elements.size()) + "]";
        }

        /**
         * Return a textual representation of this quad tree up to the specified depth.
         * @param expandDepth int; the maximum depth to expand
         * @return String; textual representation of this quad tree
         */
        public String toString(final int expandDepth)
        {
            if (expandDepth > 0)
            {
                return "SubTree [boundingBox=" + this.boundingBox + ", size=" + this.size + ", children="
                        + (this.children != null ? "[SW:" + this.children[0].toString(expandDepth - 1) + ", SE:"
                                + this.children[1].toString(expandDepth - 1) + ", NW:"
                                + this.children[2].toString(expandDepth - 1) + ", NE:"
                                + this.children[3].toString(expandDepth - 1) + "]" : "null")
                        + ", elements.size=" + this.elements.size() + "]";
            }
            else
            {
                return toString();
            }
        }

        /**
         * Dump a quad tree.
         * @param indent String; prefix for each output line
         * @return String; textual description of this quad tree.
         */
        public String dump(final String indent)
        {
            StringBuilder result = new StringBuilder();
            result.append(indent);
            result.append("SubTree [size=");
            result.append(this.size);
            result.append("] ");
            result.append(this.boundingBox);
            result.append("\n");
            String subIndent = indent + "    ";
            Iterator<RectangleAndPayload<T>> iterator = this.elements.iterator();
            for (int i = 0; i < this.elements.size(); i++)
            {
                result.append(subIndent);
                result.append(i);
                result.append(" ");
                result.append(iterator.next());
                result.append("\n");
            }
            if (this.children != null)
            {
                result.append(subIndent);
                result.append("SW");
                result.append("\n");
                result.append(this.children[0].dump(subIndent));
                result.append(subIndent);
                result.append("SE");
                result.append("\n");
                result.append(this.children[1].dump(subIndent));
                result.append(subIndent);
                result.append("NW");
                result.append("\n");
                result.append(this.children[2].dump(subIndent));
                result.append(subIndent);
                result.append("NE");
                result.append("\n");
                result.append(this.children[3].dump(subIndent));
            }
            return result.toString();
        }

    }

}

/**
 * Container for a Rectangle and a payload.
 * @param <T> Object; the payload
 */
class RectangleAndPayload<T extends Object> implements Serializable
{
    /** ... */
    private static final long serialVersionUID = 20200904L;

    /** The bounding rectangle. */
    private final Rectangle rectangle;

    /** The payload. */
    private final T payload;

    /**
     * Construct a new RectangleAndPayload object.
     * @param rectangle Rectangle; the bounding rectangle of the payload
     * @param payload T; the payload
     */
    RectangleAndPayload(final Rectangle rectangle, final T payload)
    {
        this.rectangle = rectangle;
        this.payload = payload;
    }

    /**
     * Retrieve the bounding rectangle.
     * @return Rectangle; the bounding rectangle
     */
    public Rectangle getRectangle()
    {
        return this.rectangle;
    }

    /**
     * Retrieve the payload.
     * @return T; the payload
     */
    public T getPayload()
    {
        return this.payload;
    }

    /** {@inheritDoc} */
    @Override
    public String toString()
    {
        return "RectangleAndPayload [rectangle=" + this.rectangle + ", payload=" + this.payload + "]";
    }

    /** {@inheritDoc} */
    @Override
    public int hashCode()
    {
        final int prime = 31;
        int result = 1;
        result = prime * result + ((this.payload == null) ? 0 : this.payload.hashCode());
        result = prime * result + ((this.rectangle == null) ? 0 : this.rectangle.hashCode());
        return result;
    }

    /** {@inheritDoc} */
    @Override
    public boolean equals(final Object obj)
    {
        if (this == obj)
        {
            return true;
        }
        if (obj == null)
        {
            return false;
        }
        if (getClass() != obj.getClass())
        {
            return false;
        }
        @SuppressWarnings("rawtypes")
        RectangleAndPayload other = (RectangleAndPayload) obj;
        if (this.payload == null)
        {
            if (other.payload != null)
            {
                return false;
            }
        }
        else if (!this.payload.equals(other.payload))
        {
            return false;
        }
        if (this.rectangle == null)
        {
            if (other.rectangle != null)
            {
                return false;
            }
        }
        else if (!this.rectangle.equals(other.rectangle))
        {
            return false;
        }
        return true;
    }

}