Java實現(xiàn)TFIDF算法代碼分享

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 語料路經(jīng)
  * @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();
		// 每個領(lǐng)域出現(xiàn)候選詞的文檔數(shù)
		Map<String, Integer> containsKeyMap = new HashMap<>();
		// 每個領(lǐng)域的總文檔數(shù)
		Map<String, Integer> totalDocMap = new HashMap<>();
		// TF = 候選詞出現(xiàn)次數(shù)/總詞數(shù)
		Map<String, double> tfMap = new HashMap<>();
		// scan files
		for (File f : files) {
			// 候選詞詞頻
			double termFrequency = 0;
			// 文本總詞數(shù)
			double totalTerm = 0;
			// 包含候選詞的文檔數(shù)
			int containsKeyDoc = 0;
			// 詞頻文檔計數(shù)
			int totalCount = 0;
			int fileCount = 0;
			// 標(biāo)記文件中是否出現(xiàn)候選詞
			Boolean flag = false;
			FileReader fr = new FileReader(f);
			BufferedReader br = new BufferedReader(fr);
			String s = "";
			// 計算詞頻和總詞數(shù)
			while ((s = br.readLine()) != null) {
				if (s.equals(word)) {
					termFrequency++;
					flag = true;
				}
				// 文件標(biāo)識符
				if (s.equals("$$$")) {
					if (flag) {
						containsKeyDoc++;
					}
					fileCount++;
					flag = false;
				}
				totalCount++;
			}
			// 減去文件標(biāo)識符的數(shù)量得到總詞數(shù)
			totalTerm += totalCount - fileCount;
			br.close();
			// key都為領(lǐng)域的名字
			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("該領(lǐng)域文檔數(shù)：" + fileCount);
			System.out.println("候選詞出現(xiàn)詞數(shù)：" + termFrequency);
			System.out.println("總詞數(shù)：" + totalTerm);
			System.out.println("出現(xiàn)候選詞文檔總數(shù)：" + containsKeyDoc);
			System.out.println();
		}
		//計算TF*IDF
		for (File f : files) {
			// 其他領(lǐng)域包含候選詞文檔數(shù)
			int otherContainsKeyDoc = 0;
			// 其他領(lǐng)域文檔總數(shù)
			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();
			// 計算其他領(lǐng)域包含候選詞文檔數(shù)
			for (Map.Entry<String, Integer> entry : containsKeyset) {
				if (!entry.getKey().equals(f.getName())) {
					otherContainsKeyDoc += entry.getValue();
				}
			}
			// 計算其他領(lǐng)域文檔總數(shù)
			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));
	}
}