java编写ID3决策树
说明：每个样本都会装入Data样本对象，决策树生成算法接收的是一个Array<Data>样本列表，所以构建测试数据时也要符合格式，最后生成的决策树是树的根节点，通过里面提供的showTree()方法可查看整个树结构，下面奉上源码。
Data.java
package ai.tree.data;
import java.util.HashMap;
/**
* 样本类
* @author ChenLuyang
* @date 2019/2/21
*/
public class Data implements Cloneable{
/**
* K是特征描述，V是特征值
*/
private HashMap<String,String> feature = new HashMap<String, String>();
/**
* 该样本结论
*/
private String result;
public Data(HashMap<String,String> feature,String result){
this.feature = feature;
this.result = result;
}
public HashMap<String, String> getFeature() {
return feature;
}
public String getResult() {
return result;
}
private void setFeature(HashMap<String, String> feature) {
this.feature = feature;
}
@Override
public Data clone()
{
Data object=null;
try {
object = (Data) super.clone();
object.setFeature((HashMap<String, String>) this.feature.clone());
} catch (CloneNotSupportedException e) {
e.printStackTrace();
}
return object;
}
}
DecisionTree.java
package ai.tree.algorithm;
import ai.tree.data.Data;
import java.math.BigDecimal;
import java.util.*;
/**
* @author ChenLuyang
* @date 2019/2/21
*/
public class DecisionTree {
/**
* 递归构建决策树
*
* @param dataList 样本集合
* @return ai.tree.algorithm.DecisionTree.TreeNode 使用传入样本构建的决策节点
* @author ChenLuyang
* @date 2019/2/21 16:05
*/
public TreeNode createTree(List<Data> dataList) {
//创建当前节点
TreeNode<String, String, String> nowTreeNode = new TreeNode<String, String, String>();
//当前节点的各个分支节点
Map<String, TreeNode> featureDecisionMap = new HashMap<String, TreeNode>();
//统计当前样本集中所有的分类结果
Set<String> resultSet = new HashSet<String>();
for (Data data :
dataList) {
resultSet.add(data.getResult());
}
//如果当前样本集只有一种类别，则表示不用分类了，返回当前节点
if (resultSet.size() == 1) {
String resultClassify = resultSet.iterator().next();
nowTreeNode.setResultNode(resultClassify);
return nowTreeNode;
}
//如果数据集中特征为空，则选择整个集合中出现次数最多的分类，作为分类结果
if (dataList.get(0).getFeature().size() == 0) {
Map<String, Integer> countMap = new HashMap<String, Integer>();
for (Data data :
dataList) {
Integer num = countMap.get(data.getResult());
if (num == null) {
countMap.put(data.getResult(), 1);
} else {
countMap.put(data.getResult(), num + 1);
}
}
String tmpResult = "";
Integer tmpNum = 0;
for (String res :
countMap.keySet()) {
if (countMap.get(res) > tmpNum) {
tmpNum = countMap.get(res);
tmpResult = res;
}
}
nowTreeNode.setResultNode(tmpResult);
return nowTreeNode;
}
//寻找当前最优分类
String bestLabel = chooseBestFeatureToSplit(dataList);
//提取最优特征的所有可能值
Set<String> bestLabelInfoSet = new HashSet<String>();
for (Data data :
dataList) {
bestLabelInfoSet.add(data.getFeature().get(bestLabel));
}
//使用最优特征的各个特征值进行分类
for (String labelInfo :
bestLabelInfoSet) {
for (Data data :
dataList) {
}
List<Data> branchDataList = splitDataList(dataList, bestLabel, labelInfo);
//最优特征下该特征值的节点
TreeNode branchTreeNode = createTree(branchDataList);
featureDecisionMap.put(labelInfo, branchTreeNode);
}
nowTreeNode.setDecisionNode(bestLabel, featureDecisionMap);
return nowTreeNode;
}
/**
* 计算传入数据集中的最优分类特征
*
* @param dataList
* @return int 最优分类特征的描述
* @author ChenLuyang
* @date 2019/2/21 14:12
*/
public String chooseBestFeatureToSplit(List<Data> dataList) {
//目前数据集中的特征集合
Set<String> futureSet = dataList.get(0).getFeature().keySet();
//未分类时的熵
BigDecimal baseEntropy = calcShannonEnt(dataList);
//熵差
BigDecimal bestInfoGain = new BigDecimal("0");
//最优特征
String bestFeature = "";
//按照各特征分类
for (String future :
futureSet) {
//该特征分类后的熵
BigDecimal futureEntropy = new BigDecimal("0");
//该特征的所有特征值去重集合
Set<String> futureInfoSet = new HashSet<String>();
for (Data data :
dataList) {
futureInfoSet.add(data.getFeature().get(future));
}
//按照该特征的特征值一一分类
for (String futureInfo :
futureInfoSet) {
List<Data> splitResultDataList = splitDataList(dataList, future, futureInfo);
//分类后样本数占总样本数的比例
BigDecimal tmpProb = new BigDecimal(splitResultDataList.size() + "").divide(new BigDecimal(dataList.size() + ""), 5, BigDecimal.ROUND_HALF_DOWN);
//所占比例乘以分类后的样本熵，然后再进行熵的累加
futureEntropy = futureEntropy.add(tmpProb.multiply(calcShannonEnt(splitResultDataList)));
}
BigDecimal subEntropy = baseEntropy.subtract(futureEntropy);
if (subEntropy.compareTo(bestInfoGain) >= 0) {
bestInfoGain = subEntropy;
bestFeature = future;
}
}
return bestFeature;
}
/**
* 计算传入样本集的熵值
*
* @param dataList 样本集
* @return java.math.BigDecimal 熵
* @author ChenLuyang
* @date 2019/2/22 9:41
*/
public BigDecimal calcShannonEnt(List<Data> dataList) {
//样本总数
BigDecimal sumEntries = new BigDecimal(dataList.size() + "");
//香农熵
BigDecimal shannonEnt = new BigDecimal("0");
//统计各个分类结果的样本数量
Map<String, Integer> resultCountMap = new HashMap<String, Integer>();
for (Data data :
dataList) {
Integer dataResultCount = resultCountMap.get(data.getResult());
if (dataResultCount == null) {
resultCountMap.put(data.getResult(), 1);
} else {
resultCountMap.put(data.getResult(), dataResultCount + 1);
}
}
for (String resultCountKey :
resultCountMap.keySet()) {
BigDecimal resultCountValue = new BigDecimal(resultCountMap.get(resultCountKey).toString());
BigDecimal prob = resultCountValue.divide(sumEntries, 5, BigDecimal.ROUND_HALF_DOWN);
shannonEnt = shannonEnt.subtract(prob.multiply(new BigDecimal(Math.log(prob.doubleValue()) / Math.log(2) + "")));
}
return shannonEnt;
}
/**
* 根据某个特征的特征值，进行样本数据的划分，将划分后的样本数据集返回
*
* @param dataList 待划分的样本数据集
* @param future 筛选的特征依据
* @param info 筛选的特征值依据
* @return java.util.List<ai.tree.data.Data> 按照指定特征值分类后的数据集
* @author ChenLuyang
* @date 2019/2/21 18:26
*/
public List<Data> splitDataList(List<Data> dataList, String future, String info) {
List<Data> resultDataList = new ArrayList<Data>();
for (Data data :
dataList) {
if (data.getFeature().get(future).equals(info)) {
Data newData = (Data) data.clone();
newData.getFeature().remove(future);
resultDataList.add(newData);
}
}
return resultDataList;
}
/**
* L:每一个特征的描述信息的类型
* F:特征的类型
* R:最终分类结果的类型
*/
public class TreeNode<L, F, R> {
/**
* 该节点的最优特征的描述信息
*/
private L label;
/**
* 根据不同的特征作出响应的决定。
* K为特征值，V为该特征值作出的决策节点
*/
private Map<F, TreeNode> featureDecisionMap;
/**
* 是否为最终分类节点
*/
private boolean isFinal;
/**
* 最终分类结果信息
*/
private R resultClassify;
/**
* 设置叶子节点
*
* @param resultClassify 最终分类结果
* @return void
* @author ChenLuyang
* @date 2019/2/22 18:31
*/
public void setResultNode(R resultClassify) {
this.isFinal = true;
this.resultClassify = resultClassify;
}
/**
* 设置分支节点
*
* @param label 当前分支节点的描述信息（特征）
* @param featureDecisionMap 当前分支节点的各个特征值，与其对应的子节点
* @return void
* @author ChenLuyang
* @date 2019/2/22 18:31
*/
public void setDecisionNode(L label, Map<F, TreeNode> featureDecisionMap) {
this.isFinal = false;
this.label = label;
this.featureDecisionMap = featureDecisionMap;
}
/**
* 展示当前节点的树结构
*
* @return void
* @author ChenLuyang
* @date 2019/2/22 16:54
*/
public String showTree() {
HashMap<String, String> treeMap = new HashMap<String, String>();
if (isFinal) {
String key = "result";
R value = resultClassify;
treeMap.put(key, value.toString());
} else {
String key = label.toString();
HashMap<F, String> showFutureMap = new HashMap<F, String>();
for (F f :
featureDecisionMap.keySet()) {
showFutureMap.put(f, featureDecisionMap.get(f).showTree());
}
String value = showFutureMap.toString();
treeMap.put(key, value);
}
return treeMap.toString();
}
public L getLabel() {
return label;
}
public Map<F, TreeNode> getFeatureDecisionMap() {
return featureDecisionMap;
}
public R getResultClassify() {
return resultClassify;
}
public boolean getFinal() {
return isFinal;
}
}
}
Start.java
package ai.tree.algorithm;
import ai.tree.data.Data;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
/**
* @author ChenLuyang
* @date 2019/2/22
*/
public class Start {
/**
* 构建测试样本集，测试样本如下：
样本特征：{头发长短=短发, 身材=胖, 是否戴眼镜=有眼镜} 分类：男
样本特征：{头发长短=长发, 身材=瘦, 是否戴眼镜=有眼镜} 分类：女
样本特征：{头发长短=短发, 身材=胖, 是否戴眼镜=有眼镜} 分类：女
样本特征：{头发长短=长发, 身材=胖, 是否戴眼镜=没眼镜} 分类：男
样本特征：{头发长短=短发, 身材=瘦, 是否戴眼镜=没眼镜} 分类：男
样本特征：{头发长短=长发, 身材=瘦, 是否戴眼镜=有眼镜} 分类：女
样本特征：{头发长短=长发, 身材=胖, 是否戴眼镜=有眼镜} 分类：男
* @author ChenLuyang
* @date 2019/2/21 15:34
* @return java.util.List<ai.tree.data.DecisionTreeTestData.Data> 样本集
*/
public static List<Data> createDataList(){
/**
* 样本特征描述
* @author ChenLuyang
* @date 2019/2/22 18:55
* @return java.util.List<ai.tree.data.Data>
*/
String[] labels = new String[]{"是否戴眼镜", "头发长短", "身材"};
List<Data> dataList = new ArrayList<Data>();
HashMap<String,String> feature1 = new HashMap<String, String>();
feature1.put(labels[0],"有眼镜");
feature1.put(labels[1].toString(),"短发");
feature1.put(labels[2].toString(),"胖");
dataList.add(new Data(feature1,"男"));
HashMap<String,String> feature2 = new HashMap<String, String>();
feature2.put(labels[0],"有眼镜");
feature2.put(labels[1],"长发");
feature2.put(labels[2],"瘦");
dataList.add(new Data(feature2,"女"));
HashMap<String,String> feature3 = new HashMap<String, String>();
feature3.put(labels[0],"有眼镜");
feature3.put(labels[1],"短发");
feature3.put(labels[2],"胖");
dataList.add(new Data(feature3,"女"));
HashMap<String,String> feature4 = new HashMap<String, String>();
feature4.put(labels[0],"没眼镜");
feature4.put(labels[1],"长发");
feature4.put(labels[2],"胖");
dataList.add(new Data(feature4,"男"));
HashMap<String,String> feature5 = new HashMap<String, String>();
feature5.put(labels[0],"没眼镜");
feature5.put(labels[1],"短发");
feature5.put(labels[2],"瘦");
dataList.add(new Data(feature5,"男"));
HashMap<String,String> feature6 = new HashMap<String, String>();
feature6.put(labels[0],"有眼镜");
feature6.put(labels[1],"长发");
feature6.put(labels[2],"瘦");
dataList.add(new Data(feature6,"女"));
HashMap<String,String> feature7 = new HashMap<String, String>();
feature7.put(labels[0],"有眼镜");
feature7.put(labels[1],"长发");
feature7.put(labels[2],"胖");
dataList.add(new Data(feature7,"男"));
return dataList;
}
public static void main(String[] args) {
DecisionTree decisionTree = new DecisionTree();
//使用测试样本生成决策树
DecisionTree.TreeNode tree = decisionTree.createTree(createDataList());
//展示决策树
System.out.println(tree.showTree());
}
}
生成树结构：{是否戴眼镜={没眼镜={result=男}, 有眼镜={身材={胖={头发长短={长发={result=男}, 短发={result=女}}}, 瘦={result=女}}}}}
