rbtree.java

<算法导论>第二版大部分算法实现. 1. 各类排序和顺序统计学相关 2. 数据结构 2.1 基本数据结构 2.2 散列表 2.3 二叉查找树 2.4 红黑树 2.5 数据结构

/*
 * Copyright (C) 2000-2007 Wang Pengcheng <[e-mail]>
 * Licensed to the Wang Pengcheng under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The LGPL licenses this file to You under the GNU Lesser General Public
 * Licence, Version 2.0  (the "License"); you may not use this file except in
 * compliance with the License.  You may obtain a copy of the License at
 *
 *     http://www.gnu.org/licenses/lgpl.txt
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
//14 Dec 2007
package cn.edu.whu.iss.algorithm.unit13;

import static cn.edu.whu.iss.algorithm.basictools.BasicSortTool.*;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.Set;

public class RBTree<T> implements InterfaceRBTree<T>, Set<T> {
	private RBTreeElement<T> root, sentinel;
	private int num=0;

	public RBTree() {
		sentinel = new RBTreeElement<T>(null);
		sentinel.setColor(BLACK);
		sentinel.setLeftChild(sentinel);
		sentinel.setRightChild(sentinel);
		root = sentinel;
	}

	public void leftRotate(RBTreeElement<T> x) {
		RBTreeElement<T> y = x.getRightChild();
		x.setRightChild(y.getLeftChild());
		(y.getLeftChild()).setParent(x);
		y.setParent(x.getParent());

		if (x.getParent() == sentinel) {
			root = y;
		} else if (x == x.getParent().getLeftChild()) {
			(x.getParent()).setLeftChild(y);
		} else {
			(x.getParent()).setRightChild(y);
		}

		y.setLeftChild(x);
		x.setParent(y);
	}

	public void rightRotate(RBTreeElement<T> x) {
		RBTreeElement<T> y = x.getLeftChild();
		x.setLeftChild(y.getRightChild());
		(y.getRightChild()).setParent(x);
		y.setParent(x.getParent());

		if (x.getParent() == sentinel) {
			root = y;
		} else if (x == x.getParent().getLeftChild()) {
			(x.getParent()).setLeftChild(y);
		} else {
			(x.getParent()).setRightChild(y);
		}

		y.setRightChild(x);
		x.setParent(y);

	}

	protected Object[] inorderTreeWalk() {
		return inorderTreeWalk(root);
	}

	protected Object[] inorderTreeWalk(RBTreeElement<T> root) {
		this.num = 0;
		Object[] ta = new Object[size(root)];
		inorderTreeWalk( ta, root);
		return ta;
	}

	protected void inorderTreeWalk( Object[] ob,
			RBTreeElement<T> root) {
		if (root != sentinel) {
			inorderTreeWalk(ob, root.getLeftChild());
			ob[num++] = root.getObject();
			inorderTreeWalk(ob, root.getRightChild());
		}
	}

	public T getMax() {
		return getMax(root).getObject();
	}

	private RBTreeElement<T> getMax(RBTreeElement<T> e) {
		while (e.getRightChild() != null) {
			e = e.getRightChild();
		}
		return e;
	}

	public T getMin() {
		return getMin(root).getObject();
	}

	private RBTreeElement<T> getMin(RBTreeElement<T> e) {
		while (e.getLeftChild() != null) {
			e = e.getLeftChild();
		}
		return e;
	}

	public T getPredecessor(T o) {
		RBTreeElement<T> e = searchNode(root, o);
		if (e != null) {
			return getPredecessor(e).getObject();
		} else {
			return null;
		}
	}

	public T getSuccessor(T o) {
		RBTreeElement<T> e = searchNode(root, o);
		if (e != null) {
			return getSuccessor(e).getObject();
		} else {
			return null;
		}
	}

	/*
	 * private void rbInsertFixup(RBTreeElement<T> z) { while
	 * (z.getParent().isRed()) { if (z.getParent() ==
	 * z.getParent().getParent().getLeftChild()) { RBTreeElement<T> y =
	 * z.getParent().getParent().getRightChild(); if (y.isRed()) {
	 * z.getParent().setColor(BLACK); y.setColor(BLACK);
	 * z.getParent().getParent().setColor(RED); z = z.getParent().getParent(); }
	 * else { if (z == z.getParent().getRightChild()) { z = z.getParent();
	 * leftRotate(z); } z.getParent().setColor(BLACK);
	 * z.getParent().getParent().setColor(RED);
	 * rightRotate(z.getParent().getParent()); } }else{
	 *  } } root.setColor(BLACK); }
	 */

	public T search(T key) {
		RBTreeElement<T> x = searchNode(root, key);
		if (x == null) {
			return null;
		}
		return x.getObject();
	}

	private RBTreeElement<T> getPredecessor(RBTreeElement<T> x) {
		if (x.getLeftChild() != null) {
			return getMax(x.getLeftChild());
		} else {
			RBTreeElement<T> y = x.getParent();
			while ((y != null) && (x == y.getLeftChild())) {
				x = y;
				y = x.getParent();
			}
			return y;
		}
	}

	protected RBTreeElement<T> inorderSuccessor(RBTreeElement<T> node) {
		RBTreeElement<T> des = node.getRightChild();
		while (des.getLeftChild() != sentinel) {
			des = des.getLeftChild();
		}
		return des;
	}

	private RBTreeElement<T> searchNode(RBTreeElement<T> x, T key) {
		while ((x != null) && (compare(key, x.getObject()) != 0)) {
			if (compare(key, x.getObject()) < 0) {
				x = x.getLeftChild();
			} else {
				x = x.getRightChild();
			}
		}
		return x;
	}

	private RBTreeElement<T> getSuccessor(RBTreeElement<T> x) {
		if (x.getRightChild() != null) {
			return getMin(x.getRightChild());
		} else {
			RBTreeElement<T> y = x.getParent();
			while ((y != null) && (x == y.getRightChild())) {
				x = y;
				y = x.getParent();
			}
			return y;
		}
	}

	public boolean add(T e) {
		RBTreeElement<T> targetNode = new RBTreeElement<T>(e, sentinel);
		RBTreeElement<T> parent = sentinel;
		RBTreeElement<T> node = root;
		int com = 0;
		while (node != sentinel) {
			parent = node;
			com = compareNode(targetNode, node);
			if (com == 0) {
				return false;
			}
			node = node.getChild(com);
		}
		linkParentAndChild(parent, targetNode, com);
		if (parent == sentinel) {
			root = targetNode;
		}
		repairAfterInsertion(targetNode);
		root.setParent(sentinel);
		return true;
	}

	public boolean addAll(Collection<? extends T> c) {
		Iterator<?> it = c.iterator();
		boolean f = true;
		while (it.hasNext()) {
			f = add((T) it.next());
			if (!f) {
				break;
			}
		}
		return f;
	}

	public void clear() {
		root = sentinel;
	}

	public boolean contains(Object o) {
		RBTreeElement<T> node = root;
		while (node != sentinel) {
			int com = compare(o, node.getObject());
			if (com == 0) {
				return true;
			} else {
				node = node.getChild(com);
			}
		}
		return false;
	}

	public boolean containsAll(Collection<?> c) {
		boolean f = true;
		Iterator<?> it = c.iterator();
		while (it.hasNext()) {
			f = contains(it.next());
			if (!f) {
				break;
			}
		}
		return f;
	}

	public boolean isEmpty() {
		return size() == 0;
	}

	public Iterator<T> iterator() {
		Object[] o = toArray();
		List<T> list = new ArrayList<T>();
		for(Object s:o){
			list.add((T) s);
		}
		return list.iterator();
	}

	public boolean remove(Object o) {
		T target = (T) o;
		RBTreeElement<T> node = root;
		while (node != sentinel) {
			int com = compare(target, node.getObject());
			if (com == 0) {
				if ((node.getLeftChild() == sentinel)
						|| (node.getRightChild() == sentinel)) {
					spliceOut(node);
				} else {
					RBTreeElement<T> successor = inorderSuccessor(node);
					node.setObject(successor.getObject());
					spliceOut(successor);
				}
				return true;
			} else {
				node = node.getChild(com);
			}
		}
		root.setParent(sentinel);
		return true;
	}

	protected void spliceOut(RBTreeElement<T> node) {
		RBTreeElement<T> child;
		if (node.getLeftChild() != sentinel) {
			child = node.getLeftChild();
		} else {
			child = node.getRightChild();
		}
		linkParentAndChild(node.getParent(), child, compareNode(node, node
				.getParent()));
		if (node == root) {
			root = child;
		}
		if (node.isBlack()) {
			repairAfterDeletion(child);
		}
	}

	protected void repairAfterDeletion(RBTreeElement<T> node) {
		while (node != root && (node.isBlack())) {
			RBTreeElement<T> parent = node.getParent();
			int com = compareNode(node, parent);
			RBTreeElement<T> sibling = parent.getChild(-com);
			if (sibling.isRed()) {
				sibling.setBlack();
				parent.setRed();
				rotate(-com, parent);
				sibling = node.getParent().getChild(-com);
			}
			if (sibling.hasTwoBlackChildren()) {
				sibling.setRed();
				node = node.getParent();
			} else {
				if (sibling.getChild(-com).isBlack()) {
					sibling.getChild(com).setBlack();
					sibling.setRed();
					rotate(com, sibling);
					sibling = parent.getChild(-com);
				}
				sibling.setColor(parent.getColor());
				parent.setBlack();
				sibling.getChild(-com).setBlack();
				rotate(-com, parent);
				node = root;
			}
		}
		node.setBlack();
	}

	public boolean removeAll(Collection<?> c) {
		boolean f = true;
		Iterator<?> it = c.iterator();
		while (it.hasNext()) {
			f = remove(it.next());
			if (!f) {
				break;
			}
		}
		return f;
	}

	public boolean retainAll(Collection<?> c) {
		this.clear();
		return this.addAll((Collection<? extends T>) c);
	}

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

	protected int size(RBTreeElement<T> node) {
		if (node == sentinel) {
			return 0;
		} else {
			return 1 + size(node.getLeftChild()) + size(node.getRightChild());
		}
	}

	public Object[] toArray() {
		return inorderTreeWalk();
	}

	public <T> T[] toArray(T[] a) {
		RBTree<T> tree = new RBTree<T>();
		for (T t : a) {
			tree.add(t);
		}
		return (T[]) tree.toArray();
	}

	protected void linkParentAndChild(RBTreeElement<T> parent,
			RBTreeElement<T> child, int dir) {
		parent.setChild(dir, child);
		child.setParent(parent);
	}

	protected void repairAfterInsertion(RBTreeElement<T> node) {
		while (node.getParent().isRed()) {
			RBTreeElement<T> parent = node.getParent();
			RBTreeElement<T> grandparent = parent.getParent();
			RBTreeElement<T> aunt = grandparent.getChild(-compareNode(parent,
					grandparent));
			if (aunt.isRed()) {
				parent.setBlack();
				aunt.setBlack();
				grandparent.setRed();
				node = grandparent;
			} else {
				int nodeCom = compareNode(node, parent);
				int parentCom = compareNode(parent, grandparent);
				if (nodeCom != parentCom) {
					rotate(nodeCom, parent);
					node = parent;
				}
				node.getParent().setBlack();
				node.getParent().getParent().setRed();
				rotate(parentCom, node.getParent().getParent());
			}
		}
		root.setBlack();
	}

	protected void rotate(int dir, RBTreeElement<T> node) {
		RBTreeElement<T> child = node.getChild(dir);
		RBTreeElement<T> parent = node.getParent();
		if (node.getParent() == sentinel) {
			root = child;
		}
		linkParentAndChild(node, child.getChild(-dir), dir);
		linkParentAndChild(parent, child, compareNode(node, parent));
		linkParentAndChild(child, node, -dir);
	}

	protected int compareNode(RBTreeElement<T> child, RBTreeElement<T> parent) {
		if (child == parent.getLeftChild()) {
			return -1;
		} else if (child == parent.getRightChild()) {
			return 1;
		} else {
			return compare(child.getObject(), parent.getObject());
		}
	}
	
	@Override
	public String toString() {
		Object[] o = toArray();
		String s = "[";
		int i=0;
		for(;i<o.length-1;i++){
			s+=o[i].toString()+",";
		}
		if(i==o.length-1){
			s+=o[i].toString();
		}
		s+="]";
		return s;
	}

}