ID3决策树的Java实现

package DecisionTree;

import java.io.*;
import java.util.*;

public class ID3 {

//节点类
public class DTNode {
private String attribute;
private HashMap<String, DTNode> children = new HashMap<String, DTNode>();
public String getAttribute() {
return attribute;
}
public void setAttribute(String attribute) {
this.attribute = attribute;
}
public HashMap<String, DTNode> getChildren() {
return children;
}
public void setChildren(HashMap<String, DTNode> children) {
this.children = children;
}
}

private String decisionColumn;        //决定字段

public String getDecisionColumn() {
return decisionColumn;
}

public void setDecisionColumn(String decisionColumn) {
this.decisionColumn = decisionColumn;
}

//统计每个属性在集合中出现的次数
public HashMap<String, Integer> getTypeCounts(ArrayList<String> dataset) {
HashMap<String, Integer> map = new HashMap<String, Integer>();
for (int i = 0; i < dataset.size(); i++) {
String key = dataset.get(i);
if(!map.containsKey(key))
map.put(key, 1);
else
map.put(key, map.get(key)+1);
}
return map;
}

//获取key的indexlist
public ArrayList<Integer> getIndex(String key, ArrayList<String> dataset){
ArrayList<Integer> indexlist = new ArrayList<Integer>();

for(int i = 0; i < dataset.size(); i++){
if(key.equals(dataset.get(i)))
indexlist.add(Integer.valueOf(i));
}
return indexlist;
}

//根据index获取数据集
public ArrayList<String> getSubset(ArrayList<Integer> indexlist, ArrayList<String> dataset) {
ArrayList<String> subset = new ArrayList<String>();
for(Integer i : indexlist){
subset.add(dataset.get(i.intValue()));
}
return subset;
}

//计算信息熵
public double getEntropy(ArrayList<String> dataset) {
double entropy = 0;
double prob = 0;
int sum = dataset.size();
HashMap<String, Integer> map = getTypeCounts(dataset);
Iterator<String> iter = map.keySet().iterator();
while(iter.hasNext()){
String key = iter.next();
prob = (double)map.get(key).intValue()/sum;
entropy += -1*prob*Math.log10(prob)/Math.log10(2);
}
return entropy;
}

//计算已知条件下的信息熵
public double getConditionEntropy(HashMap<String, ArrayList<String>> dataset, String IndexCol) {
double entropy = 0;
double prob = 0;
int sum = dataset.get(IndexCol).size();
HashMap<String, Integer> map = getTypeCounts(dataset.get(IndexCol));
Iterator<String> iter = map.keySet().iterator();
while(iter.hasNext()){
String key = iter.next();
prob = (double)map.get(key)/sum;
entropy+=prob*getEntropy(getSubset(getIndex(key,dataset.get(IndexCol)),dataset.get(this.decisionColumn)));
}
return entropy;
}

//建立决策树
public DTNode buildDT(HashMap<String, ArrayList<String>>dataset) {

DTNode node = new DTNode();
double info_entropy = getEntropy(dataset.get(this.decisionColumn));
//递归结束条件
if(info_entropy == 0){
node.setAttribute((dataset.get(this.decisionColumn).get(0)));
return node;
}

//求出拥有最小熵数据集的column,即最大entropy gain
double max_gain = 0;            //设置默认值
double gain = 0;
String max_column="";
Iterator<String> entropy_iter = dataset.keySet().iterator();

while(entropy_iter.hasNext()){
String key = entropy_iter.next();
if(key.equals(this.decisionColumn))
continue;
gain = getEntropy(dataset.get(decisionColumn)) - getConditionEntropy(dataset,key);  //计算信息增益
if(gain > max_gain){
max_gain = gain;
max_column = key;
}
}

node.setAttribute(max_column);
ArrayList<String> ds = dataset.get(max_column);        //最小熵数据集

//生成新数据集
Iterator<String> iter = getTypeCounts(ds).keySet().iterator();
while(iter.hasNext()){
String key = iter.next();
HashMap<String, ArrayList<String>> subset = new HashMap<String, ArrayList<String>>();
DTNode childNode;
ArrayList<Integer> indexlist = getIndex(key,ds);
Iterator<String> sub_iter = dataset.keySet().iterator();
while(sub_iter.hasNext()){
String sub_key = sub_iter.next();
if(!sub_key.equals(max_column))
subset.put(sub_key, getSubset(indexlist,dataset.get(sub_key)));
}

childNode = buildDT(subset);
node.getChildren().put(key, childNode);
}

return node;
}

//输出树
public void printDT(DTNode root){

if(root == null)
return;
System.out.println(root.attribute);
if(root.getChildren() == null)
return;

Iterator<String> iter = root.getChildren().keySet().iterator();
while(iter.hasNext()){
String key = iter.next();
System.out.print(key+" ");
printDT(root.getChildren().get(key));
}
}

//读取源文件
public HashMap<String,ArrayList<String>> read(String path){
HashMap<String,ArrayList<String>> dataset = new HashMap<String,ArrayList<String>>();
try{
File file = new File(path);
if(file.isFile() && file.exists()){ //判断文件是否存在
InputStreamReader input = new InputStreamReader(new FileInputStream(file),"UTF-8");
BufferedReader read = new BufferedReader(input);
String line = null;

ArrayList<ArrayList<String>> ds = new ArrayList<ArrayList<String>>();
while((line = read.readLine()) != null){
String[] data = line.split(",");
ArrayList<String> temp = new ArrayList<String>();
for(int i = 0; i < data.length; i++)
temp.add(data[i]);
ds.add(temp);
}

for(int i = 0; i < ds.get(0).size(); i++){
ArrayList<String> newds = new ArrayList<String>();
for(int j = 0; j < ds.size(); j++){
newds.add(ds.get(j).get(i));
}
String key = newds.get(0);
newds.remove(0);
dataset.put(key,newds);
}
input.close();
}
}catch(Exception e){
e.printStackTrace();
}

return dataset;
}

public static void main(String[] args) {
ID3 tree = new ID3();
HashMap<String,ArrayList<String>> ds = tree.read("C:"+File.separator+"Users"+File.separator+"mhua005"+File.separator+
"Desktop"+File.separator+"sample.txt");
tree.setDecisionColumn("play");
ArrayList<String> attr = new ArrayList<String>();
attr.add("outlook");
attr.add("temperature");
attr.add("humidity");
attr.add("windy");
attr.add("play");
DTNode root = tree.buildDT(ds);
tree.printDT(root);
}
}

源文件内容：

outlook,temperature,humidity,windy,play
sunny,hot,high,FALSE,no
sunny,hot,high,TRUE,no
overcast,hot,high,FALSE,yes
rainy,mild,high,FALSE,yes
rainy,cool,normal,FALSE,yes
rainy,cool,normal,TRUE,no
overcast,cool,normal,TRUE,yes
sunny,mild,high,FALSE,no
sunny,cool,normal,FALSE,yes
rainy,mild,normal,FALSE,yes
sunny,mild,normal,TRUE,yes
overcast,mild,high,TRUE,yes
overcast,hot,normal,FALSE,yes
rainy,mild,high,TRUE,no