TreeSet<String>的remove()方法失效，求排查原因

我在实现LeetCode的设计食物评分系统时，遇到TreeSet的remove()方法无法正常工作的问题——执行remove操作时会提示找不到对应食物。以下是我的实现代码：

class FoodRatings {

    class SortedSetComparator implements Comparator<String> {
        public int compare(String A, String B) {
            if (foodRatingMap.get(A) == foodRatingMap.get(B)) {
                return A.compareTo(B);
            }
            return foodRatingMap.get(B).compareTo(foodRatingMap.get(A));
        }
    }

    Map<String, SortedSet<String>> foodTypeMap;
    Map<String, String> foodMap;
    Map<String, Integer> foodRatingMap;

    public FoodRatings(String[] foods, String[] cuisines, int[] ratings) {
        foodTypeMap = new HashMap<>();
        foodMap = new HashMap<>();
        foodRatingMap = new HashMap<>();

        for (int i = 0; i<foods.length; i++) {
            foodTypeMap.putIfAbsent(cuisines[i], new TreeSet<String> (new SortedSetComparator()));
            foodMap.put(foods[i], cuisines[i]);
            foodRatingMap.put(foods[i], ratings[i]);
            foodTypeMap.get(cuisines[i]).add(foods[i]);
        }
    }

    public void changeRating(String food, int newRating) {
        foodRatingMap.put(food, newRating);
        SortedSet<String> set = foodTypeMap.get(foodMap.get(food));
        if (!set.remove(food)) {
            System.out.println("Unable to find " + food);
        }
        foodTypeMap.get(foodMap.get(food)).add(food);
    }
    
    public String highestRated(String cuisine) {
        return foodTypeMap.get(cuisine).first();
    }
    
}

测试用例：

public static void main(String args[]) {

    String[] foods = new String[] {
        "czopaaeyl", "lxoozsbh", "kbaxapl"
    };
    String[] cuisines = new String[] {
        "dmnuqeatj", "dmnuqeatj", "dmnuqeatj"
    };
    int[] ratings = new int[] {
        11, 2, 15
    };

    FoodRatings obj = new MyClass().new FoodRatings(foods, cuisines, ratings);
    obj.changeRating("czopaaeyl", 12);
    String food = obj.highestRated("dmnuqeatj");
    System.out.println(food);
    obj.changeRating("kbaxapl", 8);
    food = obj.highestRated("dmnuqeatj");
    System.out.println(food);
    obj.changeRating("lxoozsbh", 5);
    food = obj.highestRated("dmnuqeatj");
    System.out.println(food);
}

问题原因 #

1. 操作顺序错误 #

你的changeRating方法先更新了食物的评分，再尝试删除集合中的元素。TreeSet是基于自定义Comparator定位元素的，当评分修改后，该食物与集合中其他元素的比较逻辑发生变化，TreeSet无法通过原规则找到元素，导致remove()失败。

2. Integer值比较错误 #

Comparator中使用==比较Integer对象，当评分数值超过Integer缓存范围（-128到127）时，==会比较对象引用而非实际数值，导致比较逻辑错误，进一步影响TreeSet的元素定位和排序。

解决方案 #

修正1：调整操作顺序 #

先从集合中删除食物（此时评分还是旧值，匹配原比较逻辑），再更新评分，最后重新添加食物到集合：

public void changeRating(String food, int newRating) {
    SortedSet<String> set = foodTypeMap.get(foodMap.get(food));
    // 先删除旧评分对应的元素
    if (!set.remove(food)) {
        System.out.println("Unable to find " + food);
    }
    // 更新评分
    foodRatingMap.put(food, newRating);
    // 添加新评分的元素
    set.add(food);
}

修正2：正确比较Integer值 #

将Comparator中的==替换为equals()方法，确保比较的是Integer的实际数值而非引用：

class SortedSetComparator implements Comparator<String> {
    public int compare(String A, String B) {
        Integer ratingA = foodRatingMap.get(A);
        Integer ratingB = foodRatingMap.get(B);
        if (ratingA.equals(ratingB)) {
            return A.compareTo(B);
        }
        return ratingB.compareTo(ratingA);
    }
}

内容的提问来源于stack exchange，提问作者Abascus