package tree;

public class AVLTreeMap<K extends Comparable<K>,V> implements Map<K,V>
{
	private Node<K,V> root;
	private final Node<K,V> NULL_NODE = new Node<K,V>(null,null,-1,-1,null,null);
	
	public static void main(String args[])
	{
		AVLTreeMap<Integer,String> map = new AVLTreeMap<Integer,String>();
		
		map.put(10,"10");
		map.put(20,"20");
		map.put(30,"20");
		map.put(5,"5");
		map.put(7,"7");
	}
	
	public AVLTreeMap()
	{
		root = null;
	}
	
	@Override
	public void put(K key,V value)
	{
		if(root == null){
			root = new Node<K,V>(key,value,0,0,NULL_NODE,NULL_NODE);
			return;
		}
		
		root = _put(root,key,value);
	}
	
	public Node<K,V> _put(Node<K,V> cur,K key,V value)
	{
		int result = cur.key.compareTo(key);
		
		if(result < 0){
			//right
			if(cur.right == NULL_NODE){
				cur.right = new Node<K,V>(key,value,0,0,NULL_NODE,NULL_NODE);
				cur.rdepth = 1;
				return cur;
			}
			
			cur.right = _put(cur.right,key,value);
			cur.rdepth = Math.max(cur.right.ldepth,cur.right.rdepth) + 1;
			int diff = cur.rdepth - cur.ldepth;
			
			if(diff > 1){
				//一重回転か二重回転か判断する。
				//以下の式は一重回転の場合は rottype が正になるはず。負の場合は二重回転、0の場合はあり得ない。
				int rottype = result * cur.right.key.compareTo(key);
				if(rottype > 0){
					return singleRotate(cur,OP_RIGHT);
				}else{
					return doubleRotate(cur,OP_RIGHTLEFT);
				}
			}
		}else if(result > 0){
			//left
			if(cur.left == NULL_NODE){
				cur.left = new Node<K,V>(key,value,0,0,NULL_NODE,NULL_NODE);
				cur.ldepth = 1;
				return cur;
			}
			
			cur.left = _put(cur.left,key,value);
			cur.ldepth = Math.max(cur.left.ldepth,cur.left.rdepth) + 1;
			int diff = cur.ldepth - cur.rdepth;
			
			if(diff > 1){
				//一重回転か二重回転か判断する。
				//以下の式は一重回転の場合は rottype が正になるはず。負の場合は二重回転、0の場合はあり得ない。
				int rottype = result * cur.left.key.compareTo(key);
				if(rottype > 0){
					return singleRotate(cur,OP_LEFT);
				}else{
					return doubleRotate(cur,OP_LEFTRIGHT);
				}
			}
		}else{
			//find
			cur.value = value;
		}
		
		return cur;
	}
	
	public V remove(K key)
	{
		RemoveResult<K,V> r = findRemoveTarget(root,key);
		root = r.cur;
		return r.target.value;
	}
	
	private RemoveResult<K,V> findRemoveTarget(Node<K,V> cur,K key)
	{
		int result = cur.key.compareTo(key);
		if(result < 0){
			//right
			if(cur == NULL_NODE){
				return new RemoveResult<K,V>(cur,NULL_NODE);
			}
			
			RemoveResult<K,V> r = findRemoveTarget(cur.right,key);
			cur.right = r.cur;
			cur.rdepth = Math.max(cur.right.rdepth,cur.right.ldepth) + 1;
			
			int diff = cur.rdepth - cur.ldepth;
			if(diff > 1){
				int rottype = result * cur.right.key.compareTo(key);
				if(rottype > 0){
					r.cur = singleRotate(cur,OP_RIGHT);
					return r;
				}else{
					r.cur = doubleRotate(cur,OP_RIGHTLEFT);
					return r;
				}
			}
			
			r.cur = cur;
			return r;
		}else if(result > 0 ){
			//left
			if(cur == NULL_NODE){
				return new RemoveResult<K,V>(cur,NULL_NODE);
			}
			
			RemoveResult<K,V> r = findRemoveTarget(cur.left,key);
			cur.left = r.cur;
			cur.ldepth = Math.max(cur.left.rdepth,cur.left.ldepth) + 1;
			
			int diff = cur.rdepth - cur.ldepth;
			if(diff > 1){
				int rottype = result * cur.left.key.compareTo(key);
				if(rottype > 0){
					r.cur = singleRotate(cur,OP_LEFT);
					return r;
				}else{
					r.cur = doubleRotate(cur,OP_LEFTRIGHT);
					return r;
				}
			}
			
			r.cur = cur;
			return r;
		}else{
			//find
			if(cur.left == NULL_NODE && cur.right == NULL_NODE){
				//delete
				return new RemoveResult<K,V>(NULL_NODE,cur);
			}
			
			if(cur.left != NULL_NODE && cur.right == NULL_NODE){
				//replace left
				return new RemoveResult<K,V>(cur.left,cur);
			}
			
			if(cur.left == NULL_NODE && cur.right != NULL_NODE){
				//replace right
				return new RemoveResult<K,V>(cur.right,cur);
			}
			
			//both nodes are not null
			RemoveResult<K,V> r = findSubstituteNode(cur.left);
			cur.key = r.target.key;
			cur.value = r.target.value;
			cur.right = r.cur;
			cur.rdepth = Math.max(cur.ldepth,cur.rdepth) + 1;
			
			int diff = cur.rdepth - cur.ldepth;
			if(diff > 1){
				r.cur = doubleRotate(cur,OP_LEFTRIGHT);
			}
			
			return r;
		}
	}
	
	private RemoveResult<K,V> findSubstituteNode(Node<K,V> cur)
	{
		if(cur.right != NULL_NODE){
			RemoveResult<K,V> result = findSubstituteNode(cur.right);
			cur.right = result.cur;
			cur.rdepth = Math.max(cur.right.rdepth,cur.right.ldepth) + 1;
			
			int diff = cur.ldepth - cur.rdepth;
			if(diff > 1){
				result.cur = singleRotate(cur.left,OP_LEFT);
				return result;
			}
			
			result.cur = cur;
			return result;
		}
		
		RemoveResult<K,V> result = new RemoveResult<K,V>(cur.left,cur);
		return result;
	}
	
	private static class RemoveResult<K,V>
	{
		Node<K,V> cur;
		Node<K,V> target;
		
		public RemoveResult(Node<K,V> cur,Node<K,V> target)
		{
			this.cur = cur;
			this.target = target;
		}
	}
	
	private static final int OP_LEFT = 1;
	private static final int OP_RIGHT = 2;
	
	public Node<K,V> singleRotate(Node<K,V> target,int type)
	{
		switch(type){
			case OP_LEFT:
				Node<K,V> left = target.left;
				target.left = left.right;
				left.right = target;
				
				target.ldepth = Math.max(target.left.rdepth,target.left.ldepth) + 1;
				left.rdepth = Math.max(left.right.rdepth,left.right.ldepth) + 1;
				return left;
			case OP_RIGHT:
				Node<K,V> right = target.right;
				target.right = right.left;
				right.left = target;
				
				target.rdepth = Math.max(target.right.rdepth,target.right.ldepth) + 1;
				right.ldepth = Math.max(right.left.ldepth,right.left.rdepth) + 1;
				return right;
			default:
				throw new IllegalArgumentException("invalid rotation type"); 
		}
	}
	
	private static final int OP_LEFTRIGHT = 1;
	private static final int OP_RIGHTLEFT = 2;
	
	public Node<K,V> doubleRotate(Node<K,V> target,int type)
	{
		switch(type){
			case OP_LEFTRIGHT:
				target.left = singleRotate(target.left,OP_RIGHT);
				return singleRotate(target,OP_LEFT);
			case OP_RIGHTLEFT:
				target.right = singleRotate(target.right,OP_LEFT);
				return singleRotate(target,OP_RIGHT);
			default:
				throw new IllegalArgumentException("invalid rotation type.");
		}
	}
	
	@Override
	public V get(K key)
	{
		Node<K,V> cur = root;
		
		while(cur != NULL_NODE){
			int result = cur.key.compareTo(key);
			if(result < 0){
				cur = cur.right;
			}else if(result > 0){
				cur = cur.left;
			}else{
				//find
				break;
			}
		}
		
		return cur.value;
	}
	
	private static class Node<K,V>
	{
		K key;
		V value;
		
		int rdepth;
		int ldepth;
		
		Node<K,V> left;
		Node<K,V> right;
		
		public Node(K key,V value,int rdepth,int ldepth,Node<K,V> left,Node<K,V> right)
		{
			this.right = right;
			this.left = left;
			this.ldepth = ldepth;
			this.rdepth = rdepth;
			this.key = key;
			this.value = value;
		}
		
		@Override
		public String toString()
		{
			return "[" + key + ":" + value + "]";
		}
	}
}