shz.model.RedBlackBSTPQueue Maven / Gradle / Ivy

Go to download
package shz.model;

import java.util.Comparator;
import java.util.Objects;

public abstract class RedBlackBSTPQueue {
    static final class Node {
        E e;
        Node left, right;
        int size = 1;
        boolean red;

        Node(E e, boolean red) {
            this.e = e;
            this.red = red;
        }
    }

    Node root;
    final Comparator comparator;

    RedBlackBSTPQueue(Comparator comparator) {
        Objects.requireNonNull(comparator);
        this.comparator = comparator;
    }

    public final int size() {
        return size(root);
    }

    private int size(Node h) {
        return h == null ? 0 : h.size;
    }

    public final boolean isEmpty() {
        return root == null || root.size == 0;
    }

    final Node put(Node h, E e) {
        if (h == null) return new Node<>(e, true);
        int cmp = comparator.compare(e, h.e);
        if (cmp <= 0) h.left = put(h.left, e);
        else h.right = put(h.right, e);
        if (isRed(h.right) && !isRed(h.left)) h = rotateLeft(h);
        if (isRed(h.left) && isRed(h.left.left)) h = rotateRight(h);
        if (isRed(h.left) && isRed(h.right)) flipColors(h);
        h.size = size(h.left) + size(h.right) + 1;
        return h;
    }

    final boolean isRed(Node h) {
        return h != null && h.red;
    }

    final Node rotateLeft(Node h) {
        Node x = h.right;
        h.right = x.left;
        x.left = h;
        x.red = h.red;
        h.red = true;
        x.size = h.size;
        h.size = 1 + size(h.left) + size(h.right);
        return x;
    }

    final Node rotateRight(Node h) {
        Node x = h.left;
        h.left = x.right;
        x.right = h;
        x.red = h.red;
        h.red = true;
        x.size = h.size;
        h.size = 1 + size(h.left) + size(h.right);
        return x;
    }

    final void flipColors(Node h) {
        h.red = true;
        if (h.left != null) h.left.red = false;
        if (h.right != null) h.right.red = false;
    }

    final Node balance(Node h) {
        if (isRed(h.right)) h = rotateLeft(h);
        if (isRed(h.right) && !isRed(h.left)) h = rotateLeft(h);
        if (isRed(h.left) && isRed(h.left.left)) h = rotateRight(h);
        if (isRed(h.left) && isRed(h.right)) flipColors(h);
        h.size = size(h.left) + size(h.right) + 1;
        return h;
    }

    public abstract void offer(E e);

    public abstract E poll();

    public abstract E peek();

    public static final class RedBlackBSTMinPQueue extends RedBlackBSTPQueue {
        private RedBlackBSTMinPQueue(Comparator comparator) {
            super(comparator);
        }

        public static  RedBlackBSTMinPQueue of(Comparator comparator) {
            return new RedBlackBSTMinPQueue<>(comparator);
        }

        @Override
        public void offer(E e) {
            Objects.requireNonNull(e);
            if (root == null) {
                root = new Node<>(e, false);
                return;
            }
            root = put(root, e);
            root.red = false;
        }

        @Override
        public E poll() {
            E result = peek();
            if (result == null) return null;
            if (!isRed(root.left) && !isRed(root.right)) root.red = true;
            root = delTop(root);
            if (root != null && root.size > 0) root.red = false;
            return result;
        }

        @Override
        public E peek() {
            if (root == null) return null;
            Node h = root;
            while (h.left != null) h = h.left;
            return h.e;
        }

        private Node delTop(Node h) {
            if (h.left == null) return h.right;
            if (!isRed(h.left) && !isRed(h.left.left)) h = moveRedLeft(h);
            h.left = delTop(h.left);
            return balance(h);
        }

        private Node moveRedLeft(Node h) {
            flipColors(h);
            if (h.right != null && isRed(h.right.left)) {
                h.right = rotateRight(h.right);
                h = rotateLeft(h);
            }
            return h;
        }
    }

    public static final class RedBlackBSTMaxPQueue extends RedBlackBSTPQueue {
        private RedBlackBSTMaxPQueue(Comparator comparator) {
            super(comparator);
        }

        public static  RedBlackBSTMaxPQueue of(Comparator comparator) {
            return new RedBlackBSTMaxPQueue<>(comparator);
        }

        @Override
        public void offer(E e) {
            Objects.requireNonNull(e);
            if (root == null) {
                root = new Node<>(e, false);
                return;
            }
            root = put(root, e);
            root.red = false;
        }

        @Override
        public E poll() {
            E result = peek();
            if (result == null) return null;
            if (!isRed(root.left) && !isRed(root.right)) root.red = true;
            root = delTop(root);
            if (root != null && root.size > 0) root.red = false;
            return result;
        }

        @Override
        public E peek() {
            if (root == null) return null;
            Node h = root;
            while (h.right != null) h = h.right;
            return h.e;
        }

        private Node delTop(Node h) {
            if (isRed(h.left)) h = rotateRight(h);
            if (h.right == null) return h.left;
            if (!isRed(h.right) && !isRed(h.right.left)) h = moveRedRight(h);
            h.right = delTop(h.right);
            return balance(h);
        }

        private Node moveRedRight(Node h) {
            flipColors(h);
            if (h.left != null && !isRed(h.left.left)) h = rotateRight(h);
            return h;
        }
    }
}