机器学习实战（代码部分）-- 决策树（Ⅱ 按照给定特征划分数据集）

3-2 按照给定特征划分数据集

        若subDataSet = splitDataSet(dataSet, axis, value)，则dataSet[axis] == value的概率为prob = len(subDataSet) / float(len(dataSet))。则之后可以通过计算prob * 香农熵(dataSet[axis] == value)计算信息增益。

关于信息增益的计算案例参考：机器学习 -- 决策树（Ⅴ决策树构造实例）

 

【案例】若有如下数据：则splitDataSet(dataSet, axis, value)表示取出第axis列中值为value的数据集，并去除掉第axis列的数据。

splitDataSet(dataSet, 0, 0)表示在dataSet中取出第0列数值为0的数据集，再去除掉0这一列。如[0, 1, 0, 1, 'no']则变为[1, 0, 1, 'no']，[0, 1, 1, 1, 'yes']变为[1, 1, 1, 'yes']。具体情况如下图：dataSet为原始数据集，retDataSet为进行splitDataSet(dataSet, 0, 0)后得到的数据集。

 

划分数据集的函数实现：

'''
    参数：
        dataSet
        axis：划分数据集的特征
        value：需要返回的特征值
'''
def splitDataSet(dataSet, axis, value):
    retDataSet = []
    for featVec in dataSet:
        # 以下条件筛选第axis列为value的数据集
        if featVec[axis] == value:
            # 以下两行代码将数据集featVec的第axis列进行了去除
            reducedFeatVec = featVec[:axis]
            reducedFeatVec.extend(featVec[axis + 1:])
            retDataSet.append(reducedFeatVec)
    return retDataSet