每日算法 — 使用java实现TF-IDF：文本检索与关键词权重计算

一、背景介绍

1.1 什么是TF-IDF

TF-IDF（Term Frequency-Inverse Document Frequency，词频-逆文档频率）是信息检索领域最经典的关键词权重计算方法。它用于评估一个词对于一个文件在一个文件集合或语料库中的重要程度。

核心思想：
– 一个词在某篇文章中出现的频率越高，它对这篇文章越重要（TF）
– 一个词在所有文章中出现的频率越高，它的区分度越低（IDF）
– 词的重要性 = TF × IDF

1.2 应用场景

TF-IDF广泛应用于：
– 搜索引擎：计算网页与查询词的相关性
– 文本分类：提取特征词用于分类算法
– 关键词提取：自动提取文章关键词
– 推荐系统：基于内容的推荐
– 文本相似度计算：余弦相似度等

二、核心算法原理

2.1 TF（词频）计算

词频（Term Frequency）表示某个词在文档中出现的频率。

原始词频：TF(t, d) = 词t在文档d中出现的次数

归一化词频（防止长文档偏向）：
TF(t, d) = 词t在文档d中出现的次数 / 文档d的总词数

对数归一化：
TF(t, d) = 1 + log(词t在文档d中出现的次数) （如果出现次数>0）

2.2 IDF（逆文档频率）计算

逆文档频率（Inverse Document Frequency）衡量一个词的普遍程度。

公式：
IDF(t) = log( 文档总数 / (包含词t的文档数 + 1) )

加1是为了避免分母为0（平滑处理）。

2.3 TF-IDF权重

公式：
TF-IDF(t, d) = TF(t, d) × IDF(t)

含义：
– TF高、IDF高 → 词对文档很重要（是关键词）
– TF高、IDF低 → 词很常见（停用词）
– TF低、IDF高 → 词少见但在本文出现少
– TF低、IDF低 → 词不重要

2.4 倒排索引

倒排索引（Inverted Index）是搜索引擎的核心数据结构：
– 正向索引：文档 → 词列表
– 倒排索引：词 → 包含该词的文档列表

倒排索引可以快速找出包含某个词的所有文档，是TF-IDF检索的基础。

三、完整Java实现

import java.util.*;
import java.util.regex.*;

/**
 * TF-IDF算法完整实现
 * 包含文本预处理、TF计算、IDF计算、倒排索引、相关性排序
 */
public class TFIDF {

    // 停用词集合（常见无意义词汇）
    private static final Set<String> STOP_WORDS = new HashSet<>(Arrays.asList(
        "的", "了", "在", "是", "我", "有", "和", "就", "不", "人", "都", "一",
        "一个", "上", "也", "很", "到", "说", "要", "去", "你", "会", "着",
        "没有", "看", "好", "自己", "这", "那", "这个", "那个", "他", "她",
        "它", "们", "而", "与", "或", "但", "如果", "因为", "所以", "之",
        "the", "a", "an", "is", "are", "was", "were", "be", "been", "being",
        "have", "has", "had", "do", "does", "did", "will", "would", "could",
        "should", "may", "might", "can", "shall", "of", "in", "on", "at",
        "to", "for", "with", "by", "from", "as", "into", "through", "during",
        "before", "after", "above", "below", "between", "out", "off", "over",
        "under", "again", "further", "then", "once", "here", "there", "when",
        "where", "why", "how", "all", "both", "each", "few", "more", "most",
        "other", "some", "such", "no", "nor", "not", "only", "own", "same",
        "so", "than", "too", "very", "just", "and", "but", "if", "or"
    ));

    // 文档类
    static class Document {
        int id;
        String title;
        String content;
        List<String> tokens; // 分词结果
        Map<String, Integer> termFreq; // 词频统计

        public Document(int id, String title, String content) {
            this.id = id;
            this.title = title;
            this.content = content;
            this.termFreq = new HashMap<>();
        }
    }

    // 搜索结果类
    static class SearchResult {
        Document doc;
        double score;

        public SearchResult(Document doc, double score) {
            this.doc = doc;
            this.score = score;
        }

        @Override
        public String toString() {
            return String.format("[%.4f] Doc%d: %s", score, doc.id, doc.title);
        }
    }

    // 文档列表
    private List<Document> documents;
    // 倒排索引：词 -> (文档ID -> 词频)
    private Map<String, Map<Integer, Integer>> invertedIndex;
    // IDF值缓存：词 -> IDF值
    private Map<String, Double> idfCache;
    // 每个文档的总词数
    private Map<Integer, Integer> docTotalTerms;

    public TFIDF() {
        this.documents = new ArrayList<>();
        this.invertedIndex = new HashMap<>();
        this.idfCache = new HashMap<>();
        this.docTotalTerms = new HashMap<>();
    }

    // ==========================================
    // 第一部分：文本预处理
    // ==========================================

    /**
     * 简单分词（中英文混合）
     * 英文按空格和标点分割，中文按单字分割（实际应用应使用专业分词器）
     */
    public List<String> tokenize(String text) {
        List<String> tokens = new ArrayList<>();
        text = text.toLowerCase();

        // 提取英文单词
        Pattern englishPattern = Pattern.compile("[a-zA-Z]+");
        Matcher matcher = englishPattern.matcher(text);
        while (matcher.find()) {
            String word = matcher.group();
            if (!STOP_WORDS.contains(word) && word.length() > 1) {
                tokens.add(word);
            }
        }

        // 提取中文字符（按2-gram切分，提高中文处理效果）
        String chineseOnly = text.replaceAll("[^\\u4e00-\\u9fa5]", "");
        for (int i = 0; i < chineseOnly.length() - 1; i++) {
            String bigram = chineseOnly.substring(i, i + 2);
            if (!STOP_WORDS.contains(bigram)) {
                tokens.add(bigram);
            }
        }

        // 单字也加入（补充）
        for (int i = 0; i < chineseOnly.length(); i++) {
            String single = String.valueOf(chineseOnly.charAt(i));
            if (!STOP_WORDS.contains(single)) {
                tokens.add(single);
            }
        }

        return tokens;
    }

    /**
     * 统计词频
     */
    public Map<String, Integer> countTermFreq(List<String> tokens) {
        Map<String, Integer> freq = new HashMap<>();
        for (String token : tokens) {
            freq.merge(token, 1, Integer::sum);
        }
        return freq;
    }

    // ==========================================
    // 第二部分：添加文档与构建索引
    // ==========================================

    /**
     * 添加文档到语料库
     */
    public void addDocument(int id, String title, String content) {
        Document doc = new Document(id, title, content);
        doc.tokens = tokenize(content);
        doc.termFreq = countTermFreq(doc.tokens);
        documents.add(doc);

        // 更新倒排索引
        for (Map.Entry<String, Integer> entry : doc.termFreq.entrySet()) {
            String term = entry.getKey();
            int freq = entry.getValue();
            invertedIndex.computeIfAbsent(term, k -> new HashMap<>()).put(id, freq);
        }

        // 记录文档总词数
        docTotalTerms.put(id, doc.tokens.size());

        // 清空IDF缓存（因为文档集合变了）
        idfCache.clear();
    }

    // ==========================================
    // 第三部分：TF-IDF计算
    // ==========================================

    /**
     * 计算词频TF（归一化版本）
     * TF(t, d) = 词t在文档d中的出现次数 / 文档d总词数
     */
    public double calculateTF(String term, int docId) {
        Map<Integer, Integer> docMap = invertedIndex.get(term);
        if (docMap == null || !docMap.containsKey(docId)) {
            return 0.0;
        }
        int termCount = docMap.get(docId);
        int totalTerms = docTotalTerms.getOrDefault(docId, 1);
        return (double) termCount / totalTerms;
    }

    /**
     * 计算逆文档频率IDF
     * IDF(t) = log( 文档总数 / (包含词t的文档数 + 1) )
     */
    public double calculateIDF(String term) {
        if (idfCache.containsKey(term)) {
            return idfCache.get(term);
        }

        int totalDocs = documents.size();
        int docFreq = invertedIndex.containsKey(term) ? invertedIndex.get(term).size() : 0;
        // 加1平滑，避免除以0
        double idf = Math.log((double) totalDocs / (docFreq + 1));

        idfCache.put(term, idf);
        return idf;
    }

    /**
     * 计算TF-IDF值
     */
    public double calculateTFIDF(String term, int docId) {
        return calculateTF(term, docId) * calculateIDF(term);
    }

    // ==========================================
    // 第四部分：搜索与排序
    // ==========================================

    /**
     * 根据查询词搜索文档，返回按TF-IDF相似度排序的结果
     */
    public List<SearchResult> search(String query) {
        List<String> queryTerms = tokenize(query);
        if (queryTerms.isEmpty()) {
            return Collections.emptyList();
        }

        // 计算每个文档的得分
        Map<Integer, Double> scores = new HashMap<>();

        for (String term : queryTerms) {
            double idf = calculateIDF(term);
            Map<Integer, Integer> docMap = invertedIndex.get(term);

            if (docMap == null) continue;

            for (Map.Entry<Integer, Integer> entry : docMap.entrySet()) {
                int docId = entry.getKey();
                double tf = calculateTF(term, docId);
                double tfidf = tf * idf;
                // 累加查询词的TF-IDF得分
                scores.merge(docId, tfidf, Double::sum);
            }
        }

        // 转换为结果列表并排序
        List<SearchResult> results = new ArrayList<>();
        for (Map.Entry<Integer, Double> entry : scores.entrySet()) {
            Document doc = getDocumentById(entry.getKey());
            if (doc != null) {
                results.add(new SearchResult(doc, entry.getValue()));
            }
        }

        // 按得分降序排序
        results.sort((a, b) -> Double.compare(b.score, a.score));

        return results;
    }

    /**
     * 根据ID获取文档
     */
    private Document getDocumentById(int id) {
        for (Document doc : documents) {
            if (doc.id == id) return doc;
        }
        return null;
    }

    /**
     * 获取文档的关键词（按TF-IDF排序）
     */
    public List<Map.Entry<String, Double>> getTopKeywords(int docId, int topN) {
        Document doc = getDocumentById(docId);
        if (doc == null) return Collections.emptyList();

        Map<String, Double> tfidfMap = new HashMap<>();
        for (String term : doc.termFreq.keySet()) {
            double tfidf = calculateTFIDF(term, docId);
            tfidfMap.put(term, tfidf);
        }

        List<Map.Entry<String, Double>> list = new ArrayList<>(tfidfMap.entrySet());
        list.sort((a, b) -> Double.compare(b.getValue(), a.getValue()));

        return list.subList(0, Math.min(topN, list.size()));
    }

    // ==========================================
    // 第五部分：余弦相似度（可选拓展）
    // ==========================================

    /**
     * 计算两个文档的余弦相似度
     * 基于TF-IDF向量
     */
    public double cosineSimilarity(int docId1, int docId2) {
        Document doc1 = getDocumentById(docId1);
        Document doc2 = getDocumentById(docId2);
        if (doc1 == null || doc2 == null) return 0.0;

        // 找出两个文档的所有词的并集
        Set<String> allTerms = new HashSet<>();
        allTerms.addAll(doc1.termFreq.keySet());
        allTerms.addAll(doc2.termFreq.keySet());

        double dotProduct = 0.0;
        double norm1 = 0.0;
        double norm2 = 0.0;

        for (String term : allTerms) {
            double tfidf1 = calculateTFIDF(term, docId1);
            double tfidf2 = calculateTFIDF(term, docId2);
            dotProduct += tfidf1 * tfidf2;
            norm1 += tfidf1 * tfidf1;
            norm2 += tfidf2 * tfidf2;
        }

        if (norm1 == 0 || norm2 == 0) return 0.0;
        return dotProduct / (Math.sqrt(norm1) * Math.sqrt(norm2));
    }

    // ==========================================
    // 主方法：演示
    // ==========================================

    public static void main(String[] args) {
        System.out.println("╔══════════════════════════════════════╗");
        System.out.println("║      TF-IDF 算法演示系统             ║");
        System.out.println("╚══════════════════════════════════════╝\n");

        TFIDF tfidf = new TFIDF();

        // 添加示例文档
        tfidf.addDocument(1, "机器学习入门",
            "机器学习是人工智能的核心技术，通过数据训练模型来实现预测和分类。" +
            "常见的机器学习算法包括线性回归、决策树、支持向量机和神经网络等。" +
            "机器学习在图像识别、自然语言处理和推荐系统中有广泛应用。");

        tfidf.addDocument(2, "深度学习进阶",
            "深度学习是机器学习的一个分支，使用多层神经网络来学习数据的层次化表示。" +
            "深度学习在图像识别、语音识别和自然语言处理领域取得了突破性进展。" +
            "常见的深度学习框架包括TensorFlow、PyTorch和Keras。");

        tfidf.addDocument(3, "数据挖掘实战",
            "数据挖掘是从大量数据中发现有价值模式的过程。" +
            "数据挖掘技术包括聚类分析、关联规则挖掘、分类和预测等。" +
            "Apriori算法和K-Means是经典的数据挖掘算法。");

        tfidf.addDocument(4, "Java编程指南",
            "Java是一种面向对象的编程语言，具有跨平台、安全性高的特点。" +
            "Java广泛应用于企业级开发、Android应用开发和大数据处理。" +
            "Java虚拟机JVM是Java实现跨平台的核心技术。");

        tfidf.addDocument(5, "搜索引擎原理",
            "搜索引擎通过爬虫抓取网页，建立倒排索引，实现快速全文检索。" +
            "TF-IDF是搜索引擎计算网页相关性的核心算法。" +
            "PageRank算法用于评估网页的重要性和权威性。");

        System.out.println("【文档集合】");
        System.out.println("----------------------------");
        for (Document doc : tfidf.documents) {
            System.out.println("Doc" + doc.id + ": " + doc.title + " (词数: " + doc.tokens.size() + ")");
        }
        System.out.println();

        // 演示1：提取文档关键词
        System.out.println("【一、文档关键词提取（Top 5 TF-IDF）】");
        System.out.println("----------------------------");
        for (Document doc : tfidf.documents) {
            System.out.println("\nDoc" + doc.id + " - " + doc.title + ":");
            List<Map.Entry<String, Double>> keywords = tfidf.getTopKeywords(doc.id, 5);
            for (int i = 0; i < keywords.size(); i++) {
                Map.Entry<String, Double> kw = keywords.get(i);
                System.out.printf("  %d. %s (%.6f)%n", i + 1, kw.getKey(), kw.getValue());
            }
        }
        System.out.println();

        // 演示2：搜索功能
        System.out.println("【二、TF-IDF 搜索演示】");
        System.out.println("----------------------------");

        String[] queries = {"机器学习", "深度学习 神经网络", "Java 编程", "数据挖掘 算法", "搜索引擎 索引"};

        for (String query : queries) {
            System.out.println("\n查询词: \"" + query + "\"");
            List<SearchResult> results = tfidf.search(query);
            if (results.isEmpty()) {
                System.out.println("  未找到相关文档");
            } else {
                for (int i = 0; i < results.size(); i++) {
                    System.out.println("  " + (i + 1) + ". " + results.get(i));
                }
            }
        }
        System.out.println();

        // 演示3：文档相似度
        System.out.println("【三、文档余弦相似度（基于TF-IDF向量）】");
        System.out.println("----------------------------");

        int[][] pairs = {{1, 2}, {1, 3}, {1, 4}, {2, 5}, {3, 5}};
        for (int[] pair : pairs) {
            double sim = tfidf.cosineSimilarity(pair[0], pair[1]);
            Document d1 = tfidf.getDocumentById(pair[0]);
            Document d2 = tfidf.getDocumentById(pair[1]);
            System.out.printf("  %s vs %s: %.4f%n", d1.title, d2.title, sim);
        }

        System.out.println("\n【四、倒排索引统计】");
        System.out.println("----------------------------");
        System.out.println("索引词数量: " + tfidf.invertedIndex.size());
        System.out.println("文档总数: " + tfidf.documents.size());
    }
}

四、复杂度分析

4.1 时间复杂度

4.2 空间复杂度

4.3 优化策略

1. 增量索引：新文档加入时增量更新索引，无需全量重建
2. 分布式索引：大规模数据使用MapReduce分布式构建索引
3. 压缩存储：对倒排表进行差值编码压缩
4. 缓存机制：缓存热门查询结果

五、应用拓展

5.1 TF-IDF的变体

5.2 实际应用案例

1. Elasticsearch：开源搜索引擎，默认使用BM25（TF-IDF的改进版）
2. Lucene：Java编写的全文检索库，TF-IDF是核心
3. 学术论文检索：PubMed、Google Scholar
4. 电商搜索：淘宝、京东的商品搜索

5.3 局限性与改进方向

• 同义词问题：TF-IDF无法识别同义词（可用WordNet或词向量解决）
• 多义词问题：同一词在不同语境含义不同
• 语序忽略：词袋模型丢失语序信息
• 语义理解不足：仅基于词频，不理解语义

改进方向：
– 结合词向量（Word2Vec、BERT）
– 使用BM25+排序学习
– 引入语义相似度计算