Java實現TFIDF算法

算法介紹

最近要做領域概念的提取，TFIDF作為一個很經典的算法可以作為其中的一步處理。

關於TFIDF算法的介紹可以參考這篇博客http://www.ruanyifeng.com/blog/2013/03/tf-idf.html。

計算公式比較簡單，如下：

 

預處理

由於需要處理的候選詞大約后3w+，並且語料文檔數有1w+，直接挨個文本遍歷的話很耗時，每個詞處理時間都要一分鍾以上。

為了縮短時間，首先進行分詞，一個詞輸出為一行方便統計，分詞工具選擇的是HanLp。

然后，將一個領域的文檔合並到一個文件中，並用“$$$”標識符分割，方便記錄文檔數。

下面是選擇的領域語料（PATH目錄下）：

代碼實現

package edu.heu.lawsoutput;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileReader;
import java.io.FileWriter;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/**
 * @ClassName: TfIdf
 * @Description: TODO
 * @author LJH
 * @date 2017年11月12日 下午3:55:15
 */

public class TfIdf {

    static final String PATH = "E:\\corpus"; // 語料庫路徑

    public static void main(String[] args) throws Exception {

        String test = "離退休人員"; // 要計算的候選詞

        computeTFIDF(PATH, test);

    }

    /**
    * @param @param path 語料路經
    * @param @param word 候選詞
    * @param @throws Exception 
    * @return void 
    */
    static void computeTFIDF(String path, String word) throws Exception {

        File fileDir = new File(path);
        File[] files = fileDir.listFiles();

        // 每個領域出現候選詞的文檔數
        Map<String, Integer> containsKeyMap = new HashMap<>();
        // 每個領域的總文檔數
        Map<String, Integer> totalDocMap = new HashMap<>();
        // TF = 候選詞出現次數/總詞數
        Map<String, Double> tfMap = new HashMap<>();

        // scan files
        for (File f : files) {

            // 候選詞詞頻
            double termFrequency = 0;
            // 文本總詞數
            double totalTerm = 0;
            // 包含候選詞的文檔數
            int containsKeyDoc = 0;
            // 詞頻文檔計數
            int totalCount = 0;
            int fileCount = 0;
            // 標記文件中是否出現候選詞
            boolean flag = false;

            FileReader fr = new FileReader(f);
            BufferedReader br = new BufferedReader(fr);
            String s = "";

            // 計算詞頻和總詞數
            while ((s = br.readLine()) != null) {
                if (s.equals(word)) {
                    termFrequency++;
                    flag = true;
                }

                // 文件標識符
                if (s.equals("$$$")) {
                    if (flag) {
                        containsKeyDoc++;
                    }
                    fileCount++;
                    flag = false;
                }
                totalCount++;
            }

            // 減去文件標識符的數量得到總詞數
            totalTerm += totalCount - fileCount;
            br.close();
            // key都為領域的名字
            containsKeyMap.put(f.getName(), containsKeyDoc);
            totalDocMap.put(f.getName(), fileCount);
            tfMap.put(f.getName(), (double) termFrequency / totalTerm);

            System.out.println("----------" + f.getName() + "----------");
            System.out.println("該領域文檔數：" + fileCount);
            System.out.println("候選詞出現詞數：" + termFrequency);
            System.out.println("總詞數：" + totalTerm);
            System.out.println("出現候選詞文檔總數：" + containsKeyDoc);
            System.out.println();
        }
        
        //計算TF*IDF
        for (File f : files) {

            // 其他領域包含候選詞文檔數
            int otherContainsKeyDoc = 0;
            // 其他領域文檔總數
            int otherTotalDoc = 0;

            double idf = 0;
            double tfidf = 0;
            System.out.println("~~~~~" + f.getName() + "~~~~~");

            Set<Map.Entry<String, Integer>> containsKeyset = containsKeyMap.entrySet();
            Set<Map.Entry<String, Integer>> totalDocset = totalDocMap.entrySet();
            Set<Map.Entry<String, Double>> tfSet = tfMap.entrySet();

            // 計算其他領域包含候選詞文檔數
            for (Map.Entry<String, Integer> entry : containsKeyset) {
                if (!entry.getKey().equals(f.getName())) {
                    otherContainsKeyDoc += entry.getValue();
                }
            }

            // 計算其他領域文檔總數
            for (Map.Entry<String, Integer> entry : totalDocset) {
                if (!entry.getKey().equals(f.getName())) {
                    otherTotalDoc += entry.getValue();
                }
            }

            // 計算idf
            idf = log((float) otherTotalDoc / (otherContainsKeyDoc + 1), 2);

            // 計算tf*idf並輸出
            for (Map.Entry<String, Double> entry : tfSet) {
                if (entry.getKey().equals(f.getName())) {
                    tfidf = (double) entry.getValue() * idf;
                    System.out.println("tfidf:" + tfidf);
                }
            }
        }
    }

    static float log(float value, float base) {
        return (float) (Math.log(value) / Math.log(base));
    }
}            

運行結果 

測試詞為“離退休人員”，中間結果如下：

 最終結果：

結論

可以看到“離退休人員”在養老保險和社保領域，tfidf值比較高，可以作為判斷是否為領域概念的一個依據。

當然TF-IDF算法雖然很經典，但還是有許多不足，不能單獨依賴其結果做出判斷。

很多論文提出了改進方法，本文只是實現了最基本的算法。

如果有其他思路和想法歡迎討論。

---

轉載請注明原文鏈接：http://www.cnblogs.com/justcooooode/p/7831157.html