[Java] TreeMap 구현

출처

ChatGPT

자바 8 소스 코드

---

TreeMap은 자바에서 정렬된 키-값 쌍을 저장하기 위해 사용되는 NavigableMap 인터페이스의 구현체이다. TreeMap은 내부적으로 레드-블랙 트리(Red-Black Tree) 구조를 사용해 키-값을 저장하며, 이를 통해 삽입, 삭제, 검색 등의 연산이 O(log n) 시간 복잡도로 이루어진다.

• 관련 글 -> [Java] NavigableMap
• 관련 글 -> [자료 구조] 레드-블랙 트리 Red-Black Tree

 

1. TreeMap의 주요 특징

정렬된 순서

TreeMap은 항상 정렬된 상태로 데이터를 저장한다. 기본적으로 키의 자연 순서(Natural Ordering)에 따라 정렬되지만, Comparator를 통해 사용자 정의 정렬을 적용할 수 있다.

• 관련 글 -> [Java] 자연 순서 Natural Ordering

레드-블랙 트리 기반

내부적을로 레드-블랙 트리라는 이진 탐색 트리의 변형을 사용해 데이터를 관리한다. 레드-블랙 트리는 삽입, 삭제 시 자동으로 균형을 맞추기 때문에 최악의 경우에도 O(log n) 성능을 보장한다.

중복된 키 허용 안 함

TreeMap은 중복된 키를 허용하지 않는다. 만약 동일한 키가 삽입되면, 기존 값을 덮어쓴다.

Null 키 불허

TreeMap은 null 키를 허용하지 않는다. 만약 null 키를 삽입하려고 하면 NullPointerException이 발생한다. 다만 null 값은 허용된다.

 

2. TreeMap의 내부 구조

TreeMap은 내부적으로 TreeMap.Entry<K, V>라는 노드를 사용해 데이터를 저장한다. 각 노드는 트리 구조에서 하나의 키-값 쌍을 나타내며, 왼쪽 자식과 오른쪽 자식, 부모 노드, 그리고 색상 정보(레드 또는 블랙)를 포함하고 있다.

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;  // 노드의 색상 (RED 또는 BLACK)

    // Entry 생성자
    Entry(K key, V value, Entry<K,V> parent) {
        this.key = key;
        this.value = value;
        this.parent = parent;
    }

    public K getKey() { return key; }
    public V getValue() { return value; }
    public V setValue(V value) {
        V oldValue = this.value;
        this.value = value;
        return oldValue;
    }
}

 

주요 필드

• key : 트리에서 키
• value : 키에 해당하는 값
• left, right : 왼쪽 및 오른쪽 자식 노드
• parent : 부모 노드
• color : 레드-블랙 트리에서 노드의 색상(빨강 또는 검정)

 

3. 레드 블랙 트리란?

• 관련 글 -> [자료 구조] 레드-블랙 트리 Red-Black Tree

레드 블랙 트리(Red-Balck Tree)는 자가 균형 이진 탐색 트리의 일정으로, 다음과 같은 규칙을 따른다.

• 관련 글 -> [자료 구조] 자가 균형 이진 탐색 트리 Self-Balancing Binary Search Tree

• 각 노드는 빨간색(RED) 또는 검정색(BLACK) 중 하나의 색을 가진다.
• 루트 노드는 항상 검은색이다
• 모든 리프 노드(NIL)는 검은색이다. 
• 빨간색 노드의 자식은 반드시 검은색이다.(즉, 두 개의 빨간색 노드가 연속해 나타날 수 없다.)
• 루트에서 각 리프 노드까지의 모든 경로는 동일한 수의 검은색 노드를 가진다.(균형 조건)

이 규칙을 통해 트리가 불균형하게 자라는 것을 방지하고, 최악의 경우에도 O(log n)의 성능을 보장한다.

 

4. TreeMap의 주요 메소드 구현

1. 삽입 연산 put() 메소드

TreeMap에서 요소를 삽입할 때는 레드-블랙 트리의 삽입 규칙을 따른다. 요소를 트리에 삽입한 후, 레드-블랙 트리의 규칙을 위반하지 않도록 트리를 재조정(회전 및 색상 변경)한다.

public V put(K key, V value) {
    Entry<K,V> t = root;
    if (t == null) {
        compare(key, key); // type (and possibly null) check

        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) {
        do {
            parent = t;
            cmp = cpr.compare(key, t.key);
            if (cmp < 0)
                t = t.left;
            else if (cmp > 0)
                t = t.right;
            else
                return t.setValue(value);
        } while (t != null);
    }
    else {
        if (key == null)
            throw new NullPointerException();
        @SuppressWarnings("unchecked")
        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
                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;
}

 

주요 동작

• 키 비교 : 주어진 키를 트리의 기존 키들과 비교한다. 비교는 자연 순서나 Comparator를 사용한다.
• 삽입 위치 결정 : 새 키의 위치가 결정되면 트리에 삽입된다.
• 트리 재조정 : 삽입 후 트리의 균형을 맞추기 위해 회전과 색상 변경 작업이 이루어진다. 이는 fixAfterInsertion() 메소드에서 수행된다.

2. 검색 연산 get() 메소드

TreeMap에서 데이터를 검색할 때는 이진 탐색을 사용한다. 주어진 키와 트리의 각 노드의 키를 비교해 검색한다.

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();
    @SuppressWarnings("unchecked")
    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;
}

 

주요 동작

• 키 비교 : 검색할 키를 트리의 각 노드의 키와 비교하며 왼쪽 또는 오른쪽 자식 노드로 이동한다.
• 키 발견 : 주어진 키를 찾으면 해당 키에 매핑된 값을 반환한다. 찾지 못할 경우 null을 반환한다.

3. 삭제 연산 remove() 메소드

TreeMap에서 노드를 삭제할 때는 트리의 구조를 유지하기 위해 여러 가지 경우를 처리해야 한다.

1. 자식이 없는 노드 : 단순히 노드를 제거
2. 하나의 자식이 있는 노드 : 자식을 부모 노드에 연결
3. 두 개의 자식이 있는 노드 : 중위 후속자(successor) 노드로 교체한 후, 후속자를 삭제

삭제 후에도 레드-블랙 트리의 규칙을 유지하기 위해 트리를 재조정한다.

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) {
        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);

    if (replacement != null) {
        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) {
        root = null;
    } else {
        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;
        }
    }
}

 

주요 동작

• 노드 삭제 : 삭제할 노드를 찾고, 자식 노드의 수에 따라 트리 구조를 유지하면서 삭제
• 트리 재조정 : 삭제 후 레드-블랙 트리의 균형을 유지하기 위해 fixAfterDeletion() 메소드에서 트리를 재조정한다.

 

5. 레드-블랙 트리의 균형 유지 (회전 및 색상 변경)

삽입과 삭제 후 트리의 균형을 맞추기 위해 회전과 색상 변경 작업이 이루어진다. TreeMap의 내부에서는 fixAfterInsertion()과 fixAfterDeletion() 메소드를 통해 레드-블랙 트리의 규칙을 유지한다.

• 회전 : 삽입 또는 삭제로 인한 트리가 불균형해지면, 왼쪽 또는 오른쪽으로 회전해 트리의 균형을 맞춘다.
• 색상 변경 : 레드-블랙 트리의 규칙에 따라 노드의 색상을 변경해 트리의 균형을 유지한다.

 

6. TreeMap의 성능

• 시간 복잡도 : TreeMap의 모든 연산(삽입, 삭제, 검색)은 O(log n) 시간 복잡도를 가진다. 이는 레드-블랙 트리의 자가 균형 특성 덕분이다.
• 메모리 사용: TreeMap은 각 노드가 추가적인 색상 정보를 가지고 있으므로 HashMap보다 메모리 사용량이 다소 높을 수 있다.

 

결론

TreeMap은 레드-블랙 트리 기반으로 구현된 정렬된 맵으로, 삽입, 삭제, 검색 등에서 O(log n) 성능을 보장합니다. 항상 정렬된 상태를 유지해야 하거나, 범위 검색이 자주 필요할 때 매우 유용한 자료구조이다.