如何使用TreeSet查找书籍中出现频率最高的k个单词？

嘿，这个动态求Top K高频单词的问题我刚好研究过！常规操作是用Trie加堆来实现，但我发现用TreeSet就能完美搞定，代码还更简洁，插入和查询的时间复杂度都是O(log n)，实用性拉满～

核心思路 #

我们需要自定义一个对象来封装单词和它的出现次数，让这个对象实现Comparable接口，同时重写equals和hashCode方法，保证同一个单词只会对应一个对象。然后把这些对象放进TreeSet里，利用TreeSet的自动排序特性，让集合始终按单词频率从高到低排列（频率相同的话可以按字典序排序，避免排序混乱）。

代码实现示例 #

首先是自定义的MyObj类：

import java.util.Objects;

class MyObj implements Comparable<MyObj> {
    String value;
    int count;

    // 构造方法，初始计数为1
    public MyObj(String value) {
        this.value = value;
        this.count = 1;
    }

    // 增加计数并返回新的计数值
    public int incrementCount() {
        return ++count;
    }

    // 重写equals：按单词内容判断是否为同一个对象
    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;
        MyObj myObj = (MyObj) o;
        return Objects.equals(value, myObj.value);
    }

    // 重写hashCode：基于单词内容生成哈希值
    @Override
    public int hashCode() {
        return Objects.hash(value);
    }

    // 重写compareTo：先按频率降序，频率相同按字典序升序
    @Override
    public int compareTo(MyObj other) {
        // 频率不同，频率高的排前面
        if (this.count != other.count) {
            return Integer.compare(other.count, this.count);
        }
        // 频率相同，字典序小的排前面（保证排序稳定性）
        return this.value.compareTo(other.value);
    }
}

然后是结合TreeSet和HashMap的业务逻辑（用HashMap做辅助是为了快速定位单词对应的对象，比直接用TreeSet查找更高效）：

import java.util.*;

public class TopKWords {
    private final TreeSet<MyObj> sortedSet;
    private final HashMap<String, MyObj> wordMap;

    public TopKWords() {
        sortedSet = new TreeSet<>();
        wordMap = new HashMap<>();
    }

    // 动态添加单词的方法
    public void addWord(String word) {
        MyObj obj = wordMap.get(word);
        if (obj != null) {
            // 先移除旧对象（因为计数变了，排序位置会变）
            sortedSet.remove(obj);
            obj.incrementCount();
            // 重新加入集合，TreeSet会自动调整位置
            sortedSet.add(obj);
        } else {
            MyObj newObj = new MyObj(word);
            wordMap.put(word, newObj);
            sortedSet.add(newObj);
        }
    }

    // 获取频率最高的前k个单词
    public List<String> getTopK(int k) {
        List<String> result = new ArrayList<>();
        Iterator<MyObj> iterator = sortedSet.iterator();
        int count = 0;
        while (iterator.hasNext() && count < k) {
            result.add(iterator.next().value);
            count++;
        }
        return result;
    }

    // 测试用例
    public static void main(String[] args) {
        TopKWords topK = new TopKWords();
        topK.addWord("apple");
        topK.addWord("banana");
        topK.addWord("apple");
        topK.addWord("orange");
        topK.addWord("banana");
        topK.addWord("banana");

        System.out.println(topK.getTopK(2)); // 输出 [banana, apple]
    }
}

方案优势 #

• 实现简洁：不用维护复杂的Trie结构或者堆的上浮下沉逻辑，依赖TreeSet的自动排序就能搞定
• 性能稳定：插入、删除、查询都是O(log n)的时间复杂度，完全满足动态添加的场景
• 排序规则灵活：可以根据需求修改compareTo方法，比如频率相同按字典序降序，或者其他自定义规则

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