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
