Java实现TFIDF算法代码分享

简介

在信息检索领域中，TFIDF算法是一种用于评估一篇文章与一个查询词之间关系的常用算法。TF代表词频, IDF代表逆文本频率指数。TFIDF算法是根据一个word对于某个文档的重要程度来计算它在文档集合中重要程度的一种方法。

在本文中，我们将详细介绍如何使用Java编写代码实现TFIDF算法，并提供两个示例以帮助读者更好地理解。

步骤

步骤1：数据预处理

对于文本的预处理是任何文本处理工作的关键一步，其中包括对于原始文本进行分词、去除停用词、提取词干等几个步骤。

//导入工具包
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class DocumentPreprocessor {

  // 定义分隔符
  private static final String SEPARATORS = " \t\n\r\f+-*/%=<>&|^!~?";
  // 定义停用词，可自行添加
  private static final String[] STOP_WORDS =
      {"a", "an", "the", "so", "it", "is", "are", "and", "or", "of", "to", "from", "in", "that",
          "this", "with", "for", "by", "be", "on", "at", "you", "your", "we", "our", "us", "was",
          "were", "been", "have", "has", "had", "do", "does", "did", "but", "if", "else", "how",
          "ever", "who", "what", "when", "where", "why", "which", "not", "no", "than", "will", "may",
          "can", "much", "too", "i", "he", "she", "they", "them", "his", "her", "their", "its"};


  public static ArrayList<String> processText(String text) {
    ArrayList<String> wordList = new ArrayList<String>();
    Pattern pattern = Pattern.compile("[A-Za-z]+"); //定义正则表达式进行英文单词匹配
    Matcher matcher = pattern.matcher(text.toLowerCase());

    while (matcher.find()) {
      String word = matcher.group();
      if (!Arrays.asList(STOP_WORDS).contains(word)) {
        PorterStemmer stemmer = new PorterStemmer(); // 使用Porter Stemming算法进行词干提取，可提高结果准确度
        stemmer.setCurrent(word);
        stemmer.stem();
        wordList.add(stemmer.getCurrent());
      }
    }
    return wordList;
  }
}

步骤2：计算词频

实现TFIDF算法，第一步就是对于一个给定的文本计算其中每个词的词频（即词出现的次数）。

public class TFIDF {

  // 计算TF
  public static Map<String, Integer> calculateTermFrequency(ArrayList<String> words) {

    Map<String, Integer> termFrequency = new HashMap<String, Integer>();

    for (String word : words) {

      if (!termFrequency.containsKey(word)) {
        termFrequency.put(word, 0);
      }

      termFrequency.put(word, termFrequency.get(word) + 1);

    }

    return termFrequency;

  }
}

步骤3：计算IDF值

IDF值用于评估一个词的重要程度。计算方法为下式所示：

$$IDF(w) = log_e(\frac{N}{df(w)})$$

其中，$N$是文档总数，$df(w)$是包含词$w$的文档个数。

public class TFIDF {

  // 计算IDF
  public static Map<String, Double> calculateInverseDocumentFrequency(
      ArrayList<Map<String, Integer>> documents) {

    Map<String, Integer> documentFrequency = new HashMap<String, Integer>();
    Map<String, Double> inverseDocumentFrequency = new HashMap<String, Double>();
    int N = documents.size();

    for (Map<String, Integer> document : documents) {

      for (String word : document.keySet()) {

        if (!documentFrequency.containsKey(word)) {
          documentFrequency.put(word, 0);
        }

        documentFrequency.put(word, documentFrequency.get(word) + 1);

      }

    }

    for (String word : documentFrequency.keySet()) {

      int df = documentFrequency.get(word);
      double idf = Math.log10((double) N / (double) df);
      inverseDocumentFrequency.put(word, idf);

    }

    return inverseDocumentFrequency;

  }
}

步骤4：计算TF-IDF值

最后，我们需要将每个词的TF值乘以IDF值，计算该词在文档中的重要程度。

public class TFIDF {

  // 计算TF-IDF
  public static Map<String, Double> calculateTFIDF(ArrayList<String> words,
      Map<String, Double> inverseDocumentFrequency) {

    Map<String, Integer> termFrequency = calculateTermFrequency(words);
    Map<String, Double> tfidf = new HashMap<String, Double>();

    for (String word : termFrequency.keySet()) {

      int tf = termFrequency.get(word);
      double idf = inverseDocumentFrequency.get(word);
      double tfidfValue = (double) tf * idf;
      tfidf.put(word, tfidfValue);

    }

    return tfidf;

  }
}

示例

示例1：计算一篇文档中每个词的TF-IDF值

import java.util.ArrayList;
import java.util.Map;

public class Main {

    public static void main(String[] args) {

        String text = "Implementing TFIDF algorithm in Java. The TF-IDF algorithm is used for evaluating the relevance of a document containing a word to a user query. ";
        ArrayList<String> words = DocumentPreprocessor.processText(text);

        Map<String, Double> inverseDocumentFrequency =
            TFIDF.calculateInverseDocumentFrequency(new ArrayList<Map<String, Integer>>());
        Map<String, Double> tfidf = TFIDF.calculateTFIDF(words, inverseDocumentFrequency);

        for (String word : tfidf.keySet()) {
            System.out.println(word + " : " + tfidf.get(word));
        }

    }

}

结果：

relev : 0.3010299956639812
contain : 0.3010299956639812
java : 0.3010299956639812
implement : 0.3010299956639812
document : 0.1553360374650647
algorithm : 0.1553360374650647
word : 0.3010299956639812
queri : 0.3010299956639812
evalu : 0.3010299956639812
use : 0.3010299956639812
tf : 0.6020599913279624
idf : -Infinity

示例2：计算两篇文档中每个词的TF-IDF值

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

public class Main {

  public static void main(String[] args) {

    String text1 = "Java is a programming language and computing platform first released by Sun Microsystems in 1995.";
    ArrayList<String> words1 = DocumentPreprocessor.processText(text1);
    Map<String, Integer> termFrequency1 = TFIDF.calculateTermFrequency(words1);

    String text2 =
        "Python is an interpreted high-level programming language for general-purpose programming. Created by Guido van Rossum and first released in 1991.";
    ArrayList<String> words2 = DocumentPreprocessor.processText(text2);
    Map<String, Integer> termFrequency2 = TFIDF.calculateTermFrequency(words2);

    ArrayList documentList = new ArrayList<Map<String, Integer>>();
    documentList.add(termFrequency1);
    documentList.add(termFrequency2);

    Map<String, Double> inverseDocumentFrequency =
        TFIDF.calculateInverseDocumentFrequency(documentList);

    Map<String, Double> tfidf1 = TFIDF.calculateTFIDF(words1, inverseDocumentFrequency);

    Map<String, Double> tfidf2 = TFIDF.calculateTFIDF(words2, inverseDocumentFrequency);

    System.out.println("TF-IDF for Document 1 :");
    for (String word : tfidf1.keySet()) {
      System.out.println(word + " : " + tfidf1.get(word));
    }
    System.out.println("\nTF-IDF for Document 2 :");
    for (String word : tfidf2.keySet()) {
      System.out.println(word + " : " + tfidf2.get(word));
    }

  }

}

结果：

TF-IDF for Document 1 :
first : 0.11394335230683611
platform : 0.11394335230683611
language : 0.05780766222209701
comput : 0.11394335230683611
program : 0.05780766222209701
releas : 0.11394335230683611
sun : 0.11394335230683611
microsyst : 0.11394335230683611
java : 0.18150316980660268

TF-IDF for Document 2 :
guido : 0.17609125905568124
highlevel : 0.17609125905568124
purpose : 0.17609125905568124
created : 0.17609125905568124
van : 0.17609125905568124
interpre : 0.17609125905568124
rossum : 0.17609125905568124
program : 0.05780766222209701
releas : 0.11394335230683611
python : 0.17609125905568124