Java 手撸基于二叉搜索树的Map

基于二叉搜索树
支持泛型和迭代器

import java.util.*;
class TreeMap<KeyType extends Comparable<KeyType>, ValueType> implements Iterable<TreeMap<KeyType,ValueType>.Node>{
	class Node{
		private KeyType key;
		private ValueType value;
		private Node lchild;
		private Node rchild;

		public Node (KeyType key, ValueType value) {
			this.key = key;
			this.value = value;
		}
		public Node(KeyType key, Node l, Node r){
			this.key = key;
			this.lchild = l;
			this.rchild = r;
		}
		public KeyType getKey(){
			return key;
		}
		public ValueType getValue(){
			return value;
		}
	}

	private int size;
	private Node root;

	public TreeMap(){
		this.size = 0;
		this.root = null;
	}

	public Iterator<Node> iterator() {
       //匿名内部类
        return new Iterator<Node>() {
        	int index = 0;
        	int len = 0;
        	
            private Object[] arr = new Object[size];
            {
            	dfs(root);
            }
        	private void dfs(Node root){
        		if (root!=null){
        			dfs(root.lchild);
        			arr[len++] = root;
        			dfs(root.rchild);
        		}
        	}
            public boolean hasNext() {return index < size;}
            @SuppressWarnings("unchecked")
            public Node next() {
            	return (Node)arr[index++];
            }
            public void remove() { // Not implemented
                throw new UnsupportedOperationException();
            }
        };
    }

	public void add(KeyType key, ValueType value) { // 添加一项key->value键值对，如果key已存在，则覆盖
		Node currNode = new Node(key, value);
		if (root == null) {
			root = new Node(key, value);
			this.size++;
		} else {
			// find right location
			Node p = root;
			while ( true ) {
				Node next = key.compareTo(p.key)<0 ? p.lchild : p.rchild;
				if (next==null || p.key.compareTo(key)==0){
					break;
				}
				p = next;
			}
			int cmp = key.compareTo(p.key);
			// System.out.println(p.key.getClass());
			// System.out.println(key);
			// System.out.println(cmp);
			if (cmp==0) {
				p.value = value;
			} else if (cmp<0) {
			 	p.lchild = currNode;
			 	this.size++;
			} else {
				p.rchild = currNode;
				this.size++;
			}
		}
	}

	public void set(KeyType key, ValueType value) { // 将key对应的值设置为value，如果没有则追加
		if (root == null) {
			add(key, value);
		} else {
			Node p = root;
			while (true) {
				Node next = p.key.compareTo(key)<0 ? p.lchild : p.rchild;
				if (next==null || p.key.compareTo(key)==0){
					break;
				}
				p = next;
			}
			if ( p.key.compareTo(key)==0 ) {
				p.value = value;
			} else {
				add(key, value);
			}
		}

	}

	public ValueType get(KeyType key) { // 获取key对应的value，没有则返回null
		if (root==null) {
			return null;
		}
		Node p = root;
		while (true) {
			Node next = key.compareTo(p.key)<0 ? p.lchild : p.rchild;
			if (next==null || key.compareTo(p.key)==0){
				break;
			}
			p = next;
		}
		if ( key.compareTo(p.key)==0 ) {
			return p.value;
		}
		return null;
	}

	public void print() {
		f(root, 1);
	}
	private void f(Node root, int deep) {
		if (root!=null) {
			f(root.rchild, deep+1);
			for ( int i=0; i<deep; i++) {
				System.out.print("----");
			}
			System.out.print(root.key);
			System.out.print("->");
			System.out.print(root.value);
			System.out.println("");
			f(root.lchild, deep+1);
		}
	}
}

public class Test {
	public static void main(String[] args) {
		// 样例一： 数字符串每个字符的个数（虽然拿字符的ASCII码当下标开数组记录方便多了...）
		TreeMap<Character, Integer> m = new TreeMap<Character, Integer>();
		String s = "hello, world!";
		for (char c : s.toCharArray()) {
			if (m.get(c)!=null) {
				m.set(c, m.get(c)+1);
			} else {
				m.add(c, 1);
			}
		}
		m.print();
		for (TreeMap<Character,Integer>.Node n : m ) {
			System.out.printf("'%c' -> %d\n", n.getKey(), n.getValue());
		}

		// 样例二： 钱的名称与其数额的对应
		TreeMap<String, Double> money = new TreeMap<String, Double>();
		money.set("一毛", 0.1);
		money.set("两毛", 0.2);
		money.set("五毛", 0.5);
		money.set("一块", 1d);
		money.set("两块", 2d);
		money.set("五块", 5d);
		money.set("十块", 10d);
		for (TreeMap<String, Double>.Node n : money) {
			System.out.printf("%s -> %f\n", n.getKey(), n.getValue());
		}
	}
}