dynamicstatisticstree.java

来自「<算法导论>第二版大部分算法实现. 1. 各类排序和顺序统计学相关」· Java 代码 · 共 377 行

/*
 * Copyright (C) 2003-2008 Wang Pengcheng <[e-mail]>
 * Permission is granted to copy, distribute and/or modify this
 * document under the terms of the GNU Free Documentation License,
 * Version 2.0 or any later version published by the Free Software Foundation;
 * with no Invariant Sections.
 * You may obtain a copy of the License at
 *   http://www.gnu.org/licenses/lgpl.txt
 */
//12 Mar 2008
package cn.edu.whu.iss.algorithm.unit14;

import static cn.edu.whu.iss.algorithm.basictools.BasicSortTool.compare;
import cn.edu.whu.iss.algorithm.unit13.InterfaceRBTree;
import cn.edu.whu.iss.algorithm.unit13.RBTreeElement;
import static cn.edu.whu.iss.algorithm.unit14.RBNode.*;

public class DynamicStatisticsTree<T> implements InterfaceRBTree<T> {
	private RBNode<T> root, sentinel;
	private int num;
	
	public DynamicStatisticsTree() {
		sentinel = new RBNode<T>(null);
		sentinel.setColor(BLACK);
		sentinel.setSize(SENTINEL_SIZE);
		sentinel.setLeftChild(sentinel);
		sentinel.setRightChild(sentinel);
		root = sentinel;
		num=0;
	}

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

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

	public boolean contains(T o) {
		return getNode(o)!=sentinel;
	}

	private RBNode<T> getNode(T o){
		RBNode<T> node = root;
		while (node != sentinel) {
			int com = compare(o, node.getObject());
			if (com == 0) {
				return node;
			} else {
				node = node.getChild(com);
			}
		}
		return node;
	}
	
	protected void linkParentAndChild(RBNode<T> parent, RBNode<T> child, int dir) {
		parent.setChild(dir, child);
		child.setParent(parent);
	}

	protected void repairAfterInsertion(RBNode<T> node) {
		while (node.getParent().isRed()) {
			RBNode<T> parent = node.getParent();
			RBNode<T> grandparent = parent.getParent();
			RBNode<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, RBNode<T> node) {
		RBNode<T> child = node.getChild(dir);
		RBNode<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);

		// *************************************************************
		// Update the size
		//child.setSize(node.getSize());
		computeSize(node);
		computeSize(child);
		computeSize(parent);
		// *************************************************************
	}

	protected int compareNode(RBNode<T> child, RBNode<T> parent) {
		if (child == parent.getLeftChild()) {
			return -1;
		} else if (child == parent.getRightChild()) {
			return 1;
		} else {
			return compare(child.getObject(), parent.getObject());
		}
	}

	public boolean add(T e) {
		// Check the Tree have the Object of e
		if (contains(e)) {
			return false;
		}
		// ***********************************
		num++;

		RBNode<T> targetNode = new RBNode<T>(e, sentinel);
		RBNode<T> parent = sentinel;
		RBNode<T> node = root;
		int com = 0;
		while (node != sentinel) {
			parent = node;
			com = compareNode(targetNode, node);
			// *********************************************************
			// At the beginning of the add, it check the T doesn't exist
			// *********************************************************
			node.increaseSize();
			//*********************************************************
			
			node = node.getChild(com);
		}
		linkParentAndChild(parent, targetNode, com);
		if (parent == sentinel) {
			root = targetNode;
		}
		repairAfterInsertion(targetNode);
		root.setParent(sentinel);
		return true;
	}

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

	protected void repairAfterDeletion(RBNode<T> node) {
		while (node != root && (node.isBlack())) {
			RBNode<T> parent = node.getParent();
			int com = compareNode(node, parent);
			RBNode<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;
			}
		}
		//**********************************************
		computeSize(node);
		//**********************************************
		node.setBlack();
	}

	protected void spliceOut(RBNode<T> node) {
		RBNode<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;
		}
		//*********************************************
		computeSize(node.getParent());
		//computeSize(node);
		//*********************************************
		
		if (node.isBlack()) {
			repairAfterDeletion(child);
		}

	}

	public boolean remove(T target) {
		// Check the Tree have the Object of e
		if (!contains(target)) {
			return false;
		}
		// ***********************************
		num--;
		RBNode<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 {
					RBNode<T> successor = inorderSuccessor(node);
					node.setObject(successor.getObject());
					spliceOut(successor); 
				}
				
				//**************************************************
				computeSize(node);
				//**************************************************
				
				return true;
			} else {
				
				//***********************************************
				node.decreaseSize();
				//***********************************************
				
				
				node = node.getChild(com);
			}
		}
		root.setParent(sentinel);
		return true;
	}

	public void leftRotate(RBTreeElement<T> x) {
		leftRotate((RBNode<T>) x);

	}

	public void rightRotate(RBTreeElement<T> x) {
		rightRotate((RBNode<T>) x);

	}

	public void leftRotate(RBNode<T> x) {
		rotate(LEFT_CHILD, x);
	}

	public void rightRotate(RBNode<T> x) {
		rotate(RIGHT_CHILD, x);
	}

	public T selectStatisticData(int i) {
		if(i<=0||i>num){
			try{
				throw new RuntimeException("The sequence must in the [1..n]");
			}catch (Exception e) {
				e.printStackTrace();
			}
			return null;
		}
		return selectStatisticData(this.root, i);
	}

	private T selectStatisticData(RBNode<T> x, int i) {
		int r = x.getLeftChild().getSize() + 1;
		if (i == r) {
			return x.getObject();
		} else if (i < r) {
			return selectStatisticData(x.getLeftChild(), i);
		} else {
			return selectStatisticData(x.getRightChild(), i - r);
		}
	}
	
	public int rank(T target){
		if(!contains(target)){
			return -1;
		}
		RBNode<T> x = getNode(target);
		int r = x.getLeftChild().getSize()+1;
		RBNode<T> y = x;
		while(y!=sentinel){
			if(y==y.getParent().getRightChild()){
				r+=y.getParent().getLeftChild().getSize()+1;
			}
			y = y.getParent();
		}
		return r;
	}
	
	public void clear(){
		root = sentinel;
	}

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

	protected Object[] inorderTreeWalk(RBNode<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 Object[] toArray() {
		return inorderTreeWalk();
	}
	
	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;
	}
	
	private void computeSize(RBNode<T> node){
		resetSentinel();
		node.computeSize();
		resetSentinel();
	}
	private void resetSentinel(){
		sentinel.setSize(SENTINEL_SIZE);
	}
	

}