TreeMap（红黑树）源码分析
1. HashMap.Entry（红黑树节点）
private static final boolean RED = false; private static final boolean BLACK = true; static final class Entry<K,V> implements Map.Entry<K,V> { K key; // 键 V value; // 值 Entry<K,V> left; // 左孩子 Entry<K,V> right; // 右孩子 Entry<K,V> parent; // 父节点 boolean color = BLACK; // 颜色 Entry(K key, V value, Entry<K,V> parent) { this.key = key; this.value = value; this.parent = parent; } ... ... }
2. 构建TreeMap
private final Comparator<? super K> comparator; // 比较器 private transient Entry<K,V> root; // 红黑树根节点 private transient int size = 0; // 红黑树节点总数 private transient int modCount = 0; // // 调用put、remove、clear...方法时：modCount++（PrivateEntryIterator相关） public TreeMap() { // 自然排序（key实现Comparable） comparator = null; } public TreeMap(Comparator<? super K> comparator) { // 指定比较器 this.comparator = comparator; } public TreeMap(Map<? extends K, ? extends V> m) { // 用其它Map构建红黑树 comparator = null; putAll(m); } public TreeMap(SortedMap<K, ? extends V> m) { // 用其它有序Map构建红黑树 comparator = m.comparator(); try { buildFromSorted(m.size(), m.entrySet().iterator(), null, null); } catch (java.io.IOException cannotHappen) { } catch (ClassNotFoundException cannotHappen) { } }
public void putAll(Map<? extends K, ? extends V> map) { int mapSize = map.size(); if (size==0 && mapSize!=0 && map instanceof SortedMap) { // map为有序Map Comparator<?> c = ((SortedMap<?,?>)map).comparator(); if (c == comparator || (c != null && c.equals(comparator))) { ++modCount; try { buildFromSorted(mapSize, map.entrySet().iterator(), null, null); } catch (java.io.IOException cannotHappen) { } catch (ClassNotFoundException cannotHappen) { } return; } } super.putAll(map); // 依次从map中取元素添加到当前红黑树中 } private void buildFromSorted(int size, Iterator<?> it, java.io.ObjectInputStream str, V defaultVal) throws java.io.IOException, ClassNotFoundException { this.size = size; root = buildFromSorted(0, 0, size-1, computeRedLevel(size), it, str, defaultVal); } private final Entry<K,V> buildFromSorted(int level, int lo, int hi, int redLevel, Iterator<?> it, java.io.ObjectInputStream str, V defaultVal) throws java.io.IOException, ClassNotFoundException { if (hi < lo) return null; int mid = (lo + hi) >>> 1; // 中置位 Entry<K,V> left = null; if (lo < mid) // 构建左子树 left = buildFromSorted(level+1, lo, mid - 1, redLevel, it, str, defaultVal); K key; V value; if (it != null) { // 从迭代器中获取元素 if (defaultVal==null) { Map.Entry<?,?> entry = (Map.Entry<?,?>)it.next(); key = (K)entry.getKey(); value = (V)entry.getValue(); } else { key = (K)it.next(); value = defaultVal; } } else { // 从输入流中获取元素 key = (K) str.readObject(); value = (defaultVal != null ? defaultVal : (V) str.readObject()); } Entry<K,V> middle = new Entry<>(key, value, null); // 中置位节点 if (level == redLevel) middle.color = RED; if (left != null) { // 存在左子树 middle.left = left; left.parent = middle; } if (mid < hi) { // 构建右子树 Entry<K,V> right = buildFromSorted(level+1, mid+1, hi, redLevel, it, str, defaultVal); middle.right = right; right.parent = middle; } return middle; } private static int computeRedLevel(int sz) { // 树中节点全满的最后一层 int level = 0; for (int m = sz - 1; m >= 0; m = m / 2 - 1) level++; return level; }
3. get
在查找过程中，采用比较器或自然顺序比较节点大小：
1‘ 指定比较器时，优先使用比较器比较节点大小
2' 未指定比较器时，待查找键类型必须实现Comparable接口
public V get(Object key) { Entry<K,V> p = getEntry(key); return (p==null ? null : p.value); } final Entry<K,V> getEntry(Object key) { if (comparator != null) return getEntryUsingComparator(key); if (key == null) throw new NullPointerException(); Comparable<? super K> k = (Comparable<? super K>) key; Entry<K,V> p = root; while (p != null) { int cmp = k.compareTo(p.key); if (cmp < 0) // 左拐 p = p.left; else if (cmp > 0) // 右拐 p = p.right; else // 找到节点 return p; } return null; } final Entry<K,V> getEntryUsingComparator(Object key) { K k = (K) key; Comparator<? super K> cpr = comparator; if (cpr != null) { Entry<K,V> p = root; while (p != null) { int cmp = cpr.compare(k, p.key); if (cmp < 0) // 左拐 p = p.left; else if (cmp > 0) // 右拐 p = p.right; else // 找到节点 return p; } } return null; }
4. ceilingEntry和floorEntry
final int compare(Object k1, Object k2) { // 优先用comparator比较节点大小 return comparator==null ? ((Comparable<? super K>)k1).compareTo((K)k2) : comparator.compare((K)k1, (K)k2); } public Map.Entry<K,V> ceilingEntry(K key) { return exportEntry(getCeilingEntry(key)); } final Entry<K,V> getCeilingEntry(K key) { Entry<K,V> p = root; while (p != null) { int cmp = compare(key, p.key); if (cmp < 0) { if (p.left != null) // 左子树存在则左拐 p = p.left; else // 左子树不存在则当前节点为ceiling return p; } else if (cmp > 0) { if (p.right != null) { // 右子树存在则右拐 p = p.right; } else { // 右子树不存在，则从当前节点往根节点走，寻找首个比当前节点大的节点（可能不存在） Entry<K,V> parent = p.parent; Entry<K,V> ch = p; while (parent != null && ch == parent.right) { ch = parent; parent = parent.parent; } return parent; } } else // key相等的节点为ceiling return p; } return null; } public Map.Entry<K,V> floorEntry(K key) { return exportEntry(getFloorEntry(key)); } final Entry<K,V> getFloorEntry(K key) { Entry<K,V> p = root; while (p != null) { int cmp = compare(key, p.key); if (cmp > 0) { if (p.right != null) // 右子树存在则右拐 p = p.right; else // 右子树不存在则当前节点为floor return p; } else if (cmp < 0) { if (p.left != null) { // 左子树存在则左拐 p = p.left; } else { // 左子树不存在，则从当前节点往根节点走，寻找首个比当前节点小的节点（可能不存在） Entry<K,V> parent = p.parent; Entry<K,V> ch = p; while (parent != null && ch == parent.left) { ch = parent; parent = parent.parent; } return parent; } } else // key相等的节点为floor return p; } return null; // 红黑树中不存在<=key的节点 }
static <K,V> Map.Entry<K,V> exportEntry(TreeMap.Entry<K,V> e) { return (e == null) ? null : new AbstractMap.SimpleImmutableEntry<>(e); }
5. containsKey和containsValue
containsKey：getEntry返回非空
containsValue：从第一个节点开始，遍历整颗红黑树
public boolean containsKey(Object key) { return getEntry(key) != null; } public boolean containsValue(Object value) { for (Entry<K,V> e = getFirstEntry(); e != null; e = successor(e)) if (valEquals(value, e.value)) return true; return false; } public Map.Entry<K,V> firstEntry() { // 最小节点 return exportEntry(getFirstEntry()); } final Entry<K,V> getFirstEntry() { Entry<K,V> p = root; if (p != null) while (p.left != null) // 一路左拐 p = p.left; return p; }
static <K,V> TreeMap.Entry<K,V> successor(Entry<K,V> t) { // 节点t的后置节点 if (t == null) return null; else if (t.right != null) { // 右子树存在，则一步右拐再一路左拐 Entry<K,V> p = t.right; while (p.left != null) p = p.left; return p; } else { // 右子树不存在，则从当前节点往根节点走，寻找首个比当前节点大的节点（可能不存在） Entry<K,V> p = t.parent; Entry<K,V> ch = t; while (p != null && ch == p.right) { ch = p; p = p.parent; } return p; } }
6. put
1' 若已存在相同键的节点，则设置节点新值
2' 若不存在相同键的节点，则插入（黑色）节点，再以该节点为支点进行平衡调整
public V put(K key, V value) { Entry<K,V> t = root; if (t == null) { // 红黑树为空 compare(key, key); // 可能comparator为空，且key所属类也未实现Comparable接口 root = new Entry<>(key, value, null); size = 1; modCount++; return null; } int cmp; Entry<K,V> parent; Comparator<? super K> cpr = comparator; if (cpr != null) { // comparator不为空 do { parent = t; cmp = cpr.compare(key, t.key); if (cmp < 0) // 左拐 t = t.left; else if (cmp > 0) // 右拐 t = t.right; else // 存在key相等的节点 return t.setValue(value); } while (t != null); } else { // comparator为空 if (key == null) throw new NullPointerException(); Comparable<? super K> k = (Comparable<? super K>) key; do { parent = t; cmp = k.compareTo(t.key); if (cmp < 0) // 左拐 t = t.left; else if (cmp > 0) // 右拐 t = t.right; else // 存在key相等的节点 return t.setValue(value); } while (t != null); } Entry<K,V> e = new Entry<>(key, value, parent); if (cmp < 0) parent.left = e; else parent.right = e; fixAfterInsertion(e); // 插入新节点后做平衡调整 size++; modCount++; return null; } private void fixAfterInsertion(Entry<K,V> x) { x.color = RED; // x不为空 && x不为根节点 && x父节点为红色 while (x != null && x != root && x.parent.color == RED) { // 当前节点x if (parentOf(x) == leftOf(parentOf(parentOf(x)))) { // x父节点是x爷爷节点的左孩子 Entry<K,V> y = rightOf(parentOf(parentOf(x))); // x叔叔节点（可能不存在） if (colorOf(y) == RED) { // x叔叔节点为红色 setColor(parentOf(x), BLACK); setColor(y, BLACK); setColor(parentOf(parentOf(x)), RED); x = parentOf(parentOf(x)); // x爷爷节点作为当前节点：continue } else { // x叔叔节点为黑色 if (x == rightOf(parentOf(x))) { // x是父节点的右孩子 x = parentOf(x); rotateLeft(x); } // 此时，x是父节点的左孩子 setColor(parentOf(x), BLACK); setColor(parentOf(parentOf(x)), RED); rotateRight(parentOf(parentOf(x))); } } else { // x父节点是x爷爷节点的右孩子 || x爷爷节点不存在 Entry<K,V> y = leftOf(parentOf(parentOf(x))); // x叔叔节点（可能不存在） if (colorOf(y) == RED) { // x叔叔节点为红色 setColor(parentOf(x), BLACK); setColor(y, BLACK); setColor(parentOf(parentOf(x)), RED); x = parentOf(parentOf(x)); // x爷爷节点作为当前节点：continue } else { // x叔叔节点为黑色 if (x == leftOf(parentOf(x))) { // x是父节点的左孩子 x = parentOf(x); rotateRight(x); } // 此时，x是父节点的右孩子 setColor(parentOf(x), BLACK); setColor(parentOf(parentOf(x)), RED); rotateLeft(parentOf(parentOf(x))); } } } root.color = BLACK; // 根节点置为黑色 } private static <K,V> boolean colorOf(Entry<K,V> p) { // 不存在的节点（如叶子节点）为黑色 return (p == null ? BLACK : p.color); } private static <K,V> void setColor(Entry<K,V> p, boolean c) { // p不存在：noop if (p != null) p.color = c; } private static <K,V> Entry<K,V> parentOf(Entry<K,V> p) { // p不存在则p父节点也不存在 return (p == null ? null: p.parent); } private static <K,V> Entry<K,V> leftOf(Entry<K,V> p) { // p不存在则p左孩子也不存在 return (p == null) ? null: p.left; } private static <K,V> Entry<K,V> rightOf(Entry<K,V> p) { // p不存在则p右孩子也不存在 return (p == null) ? null: p.right; } private void rotateLeft(Entry<K,V> p) { // 以p为支点左旋 if (p != null) { Entry<K,V> r = p.right; // p右孩子（必须存在） // 1. r左孩子 -> p右孩子 p.right = r.left; if (r.left != null) r.left.parent = p; r.parent = p.parent; if (p.parent == null) root = r; // 2. r -> 根节点 else if (p.parent.left == p) p.parent.left = r; // 2. r -> p父节点的左孩子 else p.parent.right = r; // 2. r -> p父节点的右孩子 // 3. p -> r的左孩子 r.left = p; p.parent = r; } } private void rotateRight(Entry<K,V> p) { // 以p为支点右旋 if (p != null) { Entry<K,V> l = p.left; // p左孩子（必须存在） // 1. l右孩子 -> p左孩子 p.left = l.right; if (l.right != null) l.right.parent = p; l.parent = p.parent; if (p.parent == null) root = l; // 2. l -> 根节点 else if (p.parent.right == p) p.parent.right = l; // 2. l -> p父节点的右孩子 else p.parent.left = l; // 2. l -> p父节点的左孩子 // 3. p -> r的右孩子 l.right = p; p.parent = l; } }
7. remove
1' 若未找到待删除节点，则直接返回
2' 若待删除节点左右子树非空，则将successor的键值保存到待删除节点，待删除节点 = successor
3' 若待删除节点左右子树皆空
1'' 待删除节点为黑色，以该节点为支点进行平衡调整，再删除该节点
2’‘ 待删除节点为红色，直接删除该节点
4' 若待删除节点左右子树非皆空
1'' 待删除节点为黑色，则删除该节点，再以该节点左（右）孩子为支点进行平衡调整
2'' 待删除节点为红色，直接删除该节点
public V remove(Object key) { Entry<K,V> p = getEntry(key); // 寻找节点 if (p == null) return null; V oldValue = p.value; deleteEntry(p); return oldValue; } private void deleteEntry(Entry<K,V> p) { modCount++; size--; if (p.left != null && p.right != null) { // p左右子树非空 Entry<K,V> s = successor(p); p.key = s.key; p.value = s.value; p = s; } Entry<K,V> replacement = (p.left != null ? p.left : p.right); // p的左（右）孩子 if (replacement != null) { // p存在左（右）孩子，删除p后进行平衡调整 replacement.parent = p.parent; if (p.parent == null) root = replacement; else if (p == p.parent.left) p.parent.left = replacement; else p.parent.right = replacement; p.left = p.right = p.parent = null; if (p.color == BLACK) fixAfterDeletion(replacement); } else if (p.parent == null) { // p为红黑树唯一节点 root = null; } else { // p没有孩子节点，平衡调整完后再删除p if (p.color == BLACK) fixAfterDeletion(p); if (p.parent != null) { if (p == p.parent.left) p.parent.left = null; else if (p == p.parent.right) p.parent.right = null; p.parent = null; } } } private void fixAfterDeletion(Entry<K,V> x) { // x不为根节点 && x为黑色 while (x != root && colorOf(x) == BLACK) { // 当前节点x if (x == leftOf(parentOf(x))) { // x是x父节点的左孩子 Entry<K,V> sib = rightOf(parentOf(x)); // x兄弟节点（可能不存在） if (colorOf(sib) == RED) { // x兄弟节点为红色 setColor(sib, BLACK); setColor(parentOf(x), RED); rotateLeft(parentOf(x)); sib = rightOf(parentOf(x)); } // 此时，x兄弟节点为黑色 // x兄弟节点：左孩子黑色，右孩子黑色 if (colorOf(leftOf(sib)) == BLACK && colorOf(rightOf(sib)) == BLACK) { setColor(sib, RED); x = parentOf(x); // x父节点作为当前节点：continue } else { // x兄弟节点：左孩子红色，右孩子黑色 if (colorOf(rightOf(sib)) == BLACK) { setColor(leftOf(sib), BLACK); setColor(sib, RED); rotateRight(sib); sib = rightOf(parentOf(x)); } // 此时，x兄弟节点为黑色，x兄弟节点的右孩子为红色 setColor(sib, colorOf(parentOf(x))); setColor(parentOf(x), BLACK); setColor(rightOf(sib), BLACK); rotateLeft(parentOf(x)); x = root; } } else { // x是x父节点的右孩子 Entry<K,V> sib = leftOf(parentOf(x)); // x兄弟节点（可能不存在） if (colorOf(sib) == RED) { // x兄弟节点为红色 setColor(sib, BLACK); setColor(parentOf(x), RED); rotateRight(parentOf(x)); sib = leftOf(parentOf(x)); } // 此时，x兄弟节点为黑色 // x兄弟节点：左孩子黑色，右孩子黑色 if (colorOf(rightOf(sib)) == BLACK && colorOf(leftOf(sib)) == BLACK) { setColor(sib, RED); x = parentOf(x); // x父节点作为当前节点：continue } else { // x兄弟节点：左孩子黑色，右孩子红色 if (colorOf(leftOf(sib)) == BLACK) { setColor(rightOf(sib), BLACK); setColor(sib, RED); rotateLeft(sib); sib = leftOf(parentOf(x)); } // 此时，x兄弟节点为黑色，x兄弟节点的左孩子为红色 setColor(sib, colorOf(parentOf(x))); setColor(parentOf(x), BLACK); setColor(leftOf(sib), BLACK); rotateRight(parentOf(x)); x = root; } } } setColor(x, BLACK); }
8. entrySet、keySet和values
依次返回EntrySet、KeySet、Values
1' EntrySet、KeySet继承自AbstractSet，而Values继承自AbstractCollection
2' EntrySet、KeySet、Values对应的迭代器分别为：EntryIterator、KeyIterator、ValueIterator
3' EntryIterator、KeyIterator、ValueIterator均继承自PrivateEntryIterator，依赖PrivateEntryIterator.nextEntry分别对Entry、K、V进行迭代
4' 对EntrySet、KeySet、Values，及各自迭代器，调用remove方法，都将最终调用TreeMap的deleteEntry方法删除节点
5' 对EntrySet、KeySet、Values，不能调用add方法添加元素，否则将抛出UnsupportedOperationException
private transient EntrySet entrySet; private transient KeySet<K> navigableKeySet; public Set<Map.Entry<K,V>> entrySet() { EntrySet es = entrySet; return (es != null) ? es : (entrySet = new EntrySet()); } public Set<K> keySet() { return navigableKeySet(); } public NavigableSet<K> navigableKeySet() { KeySet<K> nks = navigableKeySet; return (nks != null) ? nks : (navigableKeySet = new KeySet<>(this)); } Iterator<K> keyIterator() { return new KeyIterator(getFirstEntry()); } public Collection<V> values() { Collection<V> vs = values; return (vs != null) ? vs : (values = new Values()); } class EntrySet extends AbstractSet<Map.Entry<K,V>> { public Iterator<Map.Entry<K,V>> iterator() { return new EntryIterator(getFirstEntry()); } public boolean contains(Object o); // TreeMap.this.getEntry public boolean remove(Object o); // TreeMap.this.deleteEntry // add方法继承自AbstractCollection：throw UnsupportedOperationException ... ... } static final class KeySet<E> extends AbstractSet<E> implements NavigableSet<E> { private final NavigableMap<E, ?> m; KeySet(NavigableMap<E,?> map) { m = map; } public Iterator<E> iterator() { if (m instanceof TreeMap) return ((TreeMap<E,?>)m).keyIterator(); else return ((TreeMap.NavigableSubMap<E,?>)m).keyIterator(); } public boolean contains(Object o); // TreeMap.this.containsKey public boolean remove(Object o); // TreeMap.this.remove // add方法继承自AbstractCollection：throw UnsupportedOperationException } class Values extends AbstractCollection<V> { public Iterator<V> iterator() { return new ValueIterator(getFirstEntry()); } public boolean contains(Object o); // TreeMap.this.containsValue public boolean remove(Object o); // TreeMap.this.deleteEntry // add方法继承自AbstractCollection：throw UnsupportedOperationException ... ... }
abstract class PrivateEntryIterator<T> implements Iterator<T> { Entry<K,V> next; // 下一节点 Entry<K,V> lastReturned; // 当前节点 int expectedModCount; PrivateEntryIterator(Entry<K,V> first) { expectedModCount = modCount; lastReturned = null; next = first; } public final boolean hasNext(); // next != null final Entry<K,V> nextEntry(); // TreeMap.this.successor final Entry<K,V> prevEntry(); // TreeMap.this.predecessor public void remove(); // TreeMap.this.deleteEntry } final class EntryIterator extends PrivateEntryIterator<Map.Entry<K,V>> { EntryIterator(Entry<K,V> first) { super(first); } public Map.Entry<K,V> next() { return nextEntry(); } } final class ValueIterator extends PrivateEntryIterator<V> { ValueIterator(Entry<K,V> first) { super(first); } public V next() { return nextEntry().value; } } final class KeyIterator extends PrivateEntryIterator<K> { KeyIterator(Entry<K,V> first) { super(first); } public K next() { return nextEntry().key; } }
9. tailMap
1’ tailMap是一个设置低位键（包括低位边界），而未设置高位键的NavigableSubMap
2' 在>=低位键的范围内，可以对tailMap进行get、put、remove等操作
public SortedMap<K,V> tailMap(K fromKey) { return tailMap(fromKey, true); } public NavigableMap<K,V> tailMap(K fromKey, boolean inclusive) { return new AscendingSubMap<>(this, false, fromKey, inclusive, // 设置低位键 true, null, true); // 未设置高位键 } static final class AscendingSubMap<K,V> extends NavigableSubMap<K,V> { AscendingSubMap(TreeMap<K,V> m, boolean fromStart, K lo, boolean loInclusive, boolean toEnd, K hi, boolean hiInclusive) { super(m, fromStart, lo, loInclusive, toEnd, hi, hiInclusive); } ... ... } abstract static class NavigableSubMap<K,V> extends AbstractMap<K,V> implements NavigableMap<K,V>, java.io.Serializable { final TreeMap<K,V> m; final K lo, hi; // 高低键 final boolean fromStart, toEnd; // 是否设置高低键 final boolean loInclusive, hiInclusive; // 是否包含高低边界 NavigableSubMap(TreeMap<K,V> m, boolean fromStart, K lo, boolean loInclusive, boolean toEnd, K hi, boolean hiInclusive) { if (!fromStart && !toEnd) { if (m.compare(lo, hi) > 0) throw new IllegalArgumentException("fromKey > toKey"); } else { if (!fromStart) // type check m.compare(lo, lo); if (!toEnd) m.compare(hi, hi); } this.m = m; this.fromStart = fromStart; this.lo = lo; this.loInclusive = loInclusive; this.toEnd = toEnd; this.hi = hi; this.hiInclusive = hiInclusive; } final boolean tooLow(Object key) { if (!fromStart) { // 设置了低位键 int c = m.compare(key, lo); if (c < 0 || (c == 0 && !loInclusive)) return true; } return false; } final boolean tooHigh(Object key) { if (!toEnd) { // 设置了高位键 int c = m.compare(key, hi); if (c > 0 || (c == 0 && !hiInclusive)) return true; } return false; } final boolean inRange(Object key) { return !tooLow(key) && !tooHigh(key); // 在高低键范围内 } public final V put(K key, V value) { if (!inRange(key)) // 范围检查 throw new IllegalArgumentException("key out of range"); return m.put(key, value); } public final V get(Object key) { return !inRange(key)/*范围检查*/ ? null : m.get(key); } public final V remove(Object key) { return !inRange(key)/*范围检查*/ ? null : m.remove(key); } ... ... }
