TreeSet의 remove 메서드 관련 이슈 정리

LeetCode 문제를 풀던 와중에 TreeSet을 사용하게 되었고, TreeSet의 요소를 지우는 remove 메서드도 사용했다.

일부 테스트케이스에서 계속해서 정답이 틀리다고 나오는 것이었다.

해당 문제를 풀기위해서 조금 더 디버깅을 해보니 내가 원하는 요소가 제대로 remove 되지 않는 이슈가 있었다.

LeetCode 문제

LeetCode 문제 링크

문제에 대해서 간단하게 설명하면, 요리는 각각의 점수를 가지고 있고, 특정 메서드를 호출하면 최고 요리의 이름을 가져오는 메서드를 설계하면 되는 문제이다.

TreeSet 이슈 정리

1. Food Class를 정의한다. (문제에서 Food의 이름은 고유하다고 했기 때문에 이름만 같으면 동일하다고 판단하기 위해서 equals 와 hashCode를 override 했다.

@Getter
class Food {
    final String name;
    final int rating;

    public Food(final String name, final int rating) {
        this.name = name;
        this.rating = rating;
    }

    @Override
    public boolean equals(final Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;
        final Food food = (Food) o;
        return Objects.equals(getName(), food.getName());
    }

    @Override
    public int hashCode() {
        return Objects.hash(getName());
    }
}

2. TreeSet은 Food 라는 객체를 가지고, Food를 정렬하기 위해서 FOOD_COMPARATOR를 사용한다.

// Raiting에 따라 우선 정렬을 하고 동일한 점수일 경우 이름으로 정렬을 한다.
final Comparator<Food> FOOD_COMPARATOR = Comparator.comparingInt(Food::getRating).reversed().thenComparing(Food::getName);


final TreeSet<Food> set = new TreeSet<>(FOOD_COMPARATOR)

3. set에서 특정 객체를 지우기 위해서 Food를 생성 후 해당 Food를 기반으로 삭제를 한다.

set.remove(new Food(foodName, 0));

4. 내 예상대로 라면 new Food(foodName, 10) 과 new Food(foodName, 0)은 equals가 true이기 때문에,
   set에 new Food(foodName, 10)이 들어가 있으면 set.remove(new Food(foodName, 0)) 메서드로 new Food(foodName, 10) 객체의 삭제가 가능하다고 생각을 했다.

final TreeSet<Food> set = new TreeSet<>(FOOD_COMPARATOR)

set.add(new Food(foodName, 10));

boolean result = set.remove(new Food(foodName, 0));

//내 예상
result == true

//실제 결과 
result == false

5. 그러나 계속해서 remove를 하지 못하는 것이었다.
   그래서 TreeSet의 remove 코드를 확인하기로 했다.
   Comparotor를 넣어서 TreeSet을 생성하면 TreeMap을 사용해서 TreeSet이 만들어진다.
   remove 메서드는 TreeMap에게 삭제 로직을 위임한다.

//TreeSet
private transient NavigableMap<E,Object> m;

public TreeSet(Comparator<? super E> comparator) {
    this(new TreeMap<>(comparator));
}

TreeSet(NavigableMap<E,Object> m) {
    this.m = m;
}

public boolean remove(Object o) {
    return m.remove(o)==PRESENT;
}

6. TreeMap에서는 일단 넘어온 key의 해당 Entry를 찾고 해당 Entry를 찾으면 삭제 로직을 진행한다.

//TreeMap
public V remove(Object key) {
    Entry<K,V> p = getEntry(key);
    if (p == null)
        return null;

    V oldValue = p.value;
    deleteEntry(p);
    return oldValue;
}

7. 그래서 나는 해당 메서드의 순서대로 getEntry 로직을 우선 살펴봤고 getEntry 메서드를 확인하니 왜 내 예상과 다르게 해당 Food가 삭제가 안되었는지 알 수 있었다.

8. 코드를 살펴보면 root에서 부터 while을 돌면서 Comparator를 사용해서 compareTo메서드를 통해 key에 해당하는 Entry를 찾는 로직이 들어 있다.
   (TreeMap 같은 경우 RedBlackTree 로 구현이 되어있다.)
   

   여기서 Comparator의 비교 값을 고려하지 않은채 new Food(foodName, 0) 와 같이 넣게 되면 compareTo 메서드에서 전혀 다른 비교값이 나오게 된다.

   그리고 내가 원하는 객체를 제대로 찾아가지 못하기 때문에 해당 Entry를 찾지 못하고 remove 메서드의 return으로 false가 나오는 것이었다.

//TreeMap
final TreeMap.Entry<K,V> getEntry(Object key) {
    // Offload comparator-based version for sake of performance
    if (comparator != null)
        return getEntryUsingComparator(key);
    if (key == null)
        throw new NullPointerException();
    @SuppressWarnings("unchecked")
    Comparable<? super K> k = (Comparable<? super K>) key;
    TreeMap.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;
}

후기

자료 구조의 특성에 대해서는 이미 어느정도 알고 있다라고 생각하고 있었는데, 실제로 이슈로 마주치니 "나는 겉핥기로만 알고 있었다"라는 사실을 다시 한번 더 깨닫게 되었다.

머리로만 알고 있던 자료구조를 다시 한번 더 생각해보는 좋은 계기였다.