Java 实现决策树算法：从原理到代码实战

核心算法：信息增益计算

构建决策树的关键在于每一步选择哪个属性作为分裂点。我们采用信息增益最大化的原则，即计算每个属性划分后的加权熵值，选择熵值最小的那个属性。

在 chooseBestTestAttribute 方法中，我们遍历所有可用属性，统计各分支下的样本分布，利用公式 $H(S) = -\sum p_i \log_2 p_i$ 计算信息量。

    static Object[] chooseBestTestAttribute(Map<String, List<Sample>> categoryToSamples, String[] attrNames) {
        int minIndex = -1;
        double minValue = Double.MAX_VALUE;
        Map<String, Map<String, List<Sample>>> minSplits = null;

        for (int attrIndex = 0; attrIndex < attrNames.length; attrIndex++) {
            int allCount = 0;
            Map<String, Map<String, List<Sample>>> curSplits = new HashMap<>();

            // 按当前属性构建 Map：属性值->(分类->样本列表)
            for (Entry<String, List<Sample>> entry : categoryToSamples.entrySet()) {
                Object category = entry.getKey();
                List<Sample> samples = entry.getValue();
                for (Sample sample : samples) {
                    Object attrValue = sample.getAttribute(attrNames[attrIndex]);
                    Map<String, List<Sample>> split = curSplits.get(attrValue);
                    if (split == null) {
                        split = new HashMap<>();
                        curSplits.put(attrValue, split);
                    }
                    List<Sample> splitSamples = split.get(category);
                    if (splitSamples == null) {
                        splitSamples = new LinkedList<>();
                        split.put(category, splitSamples);
                    }
                    splitSamples.add(sample);
                }
                allCount += samples.size();
            }

            // 计算信息量之和
            double curValue = 0.0;
            for (Map<String, List<Sample>> splits : curSplits.values()) {
                double perSplitCount = 0;
                for (List<Sample> list : splits.values()) perSplitCount += list.size();

                double perSplitValue = 0.0;
                for (List<Sample> list : splits.values()) {
                    double p = list.size() / perSplitCount;
                    perSplitValue -= p * (Math.log(p) / Math.log(2));
                }
                curValue += (perSplitCount / allCount) * perSplitValue;
            }

            if (minValue > curValue) {
                minIndex = attrIndex;
                minValue = curValue;
                minSplits = curSplits;
            }
        }
        return new Object[]{minIndex, minValue, minSplits};
    }

递归构建与输出

一旦确定了最佳属性，就将样本集按该属性的不同取值切分成若干子集，递归调用 generateDecisionTree。当样本属于同一类或没有剩余属性时，递归终止。

最后通过 outputDecisionTree 将树结构以缩进形式打印出来，方便人工检查逻辑是否正确。

    static void outputDecisionTree(Object obj, int level, Object from) {
        for (int i = 0; i < level; i++) System.out.print("|-----");
        if (from != null) System.out.printf("(%s):", from);
        if (obj instanceof Tree) {
            Tree tree = (Tree) obj;
            String attrName = tree.getAttribute();
            System.out.printf("[%s = ?]\n", attrName);
            for (Object attrValue : tree.getAttributeValues()) {
                Object child = tree.getChild(attrValue);
                outputDecisionTree(child, level + 1, attrName + " = " + attrValue);
            }
        } else {
            System.out.printf("[CATEGORY = %s]\n", obj);
        }
    }

完整代码整合

将上述片段整合到一个类中即可运行。注意 readSamples 方法中硬编码了示例数据，实际应用中可替换为文件读取。

    public static void main(String[] args) throws Exception {
        String[] attrNames = new String[] { "AGE", "INCOME", "STUDENT", "CREDIT_RATING" };
        Map<String, List<Sample>> samples = readSamples(attrNames);
        Object decisionTree = generateDecisionTree(samples, attrNames);
        outputDecisionTree(decisionTree, 0, null);
    }

    static Map<String, List<Sample>> readSamples(String[] attrNames) {
        Object[][] rawData = new Object[][] {
            { "<30 ", "High ", "No ", "Fair ", "0" },
            { "<30 ", "High ", "No ", "Excellent", "0" },
            { "30-40", "High ", "No ", "Fair ", "1" },
            { ">40 ", "Medium", "No ", "Fair ", "1" },
            { ">40 ", "Low ", "Yes", "Fair ", "1" },
            { ">40 ", "Low ", "Yes", "Excellent", "0" },
            { "30-40", "Low ", "Yes", "Excellent", "1" },
            { "<30 ", "Medium", "No ", "Fair ", "0" },
            { "<30 ", "Low ", "Yes", "Fair ", "1" },
            { ">40 ", "Medium", "Yes", "Fair ", "1" },
            { "<30 ", "Medium", "Yes", "Excellent", "1" },
            { "30-40", "Medium", "No ", "Excellent", "1" },
            { "30-40", "High ", "Yes", "Fair ", "1" },
            { ">40 ", "Medium", "No ", "Excellent", "0" }
        };
        Map<String, List<Sample>> ret = new HashMap<>();
        for (Object[] row : rawData) {
            Sample sample = new Sample();
            int n = row.length;
            for (int i = 0; i < n - 1; i++) sample.setAttribute(attrNames[i], row[i]);
            sample.setCategory(row[n - 1]);
            List<Sample> list = ret.get(row[n - 1]);
            if (list == null) {
                list = new LinkedList<>();
                ret.put(row[n - 1], list);
            }
            list.add(sample);
        }
        return ret;
    }

    static Object generateDecisionTree(Map<String, List<Sample>> categoryToSamples, String[] attrNames) {
        if (categoryToSamples.size() == 1) return categoryToSamples.keySet().iterator().next();
        if (attrNames.length == 0) {
            int max = 0;
            Object maxCategory = null;
            for (Entry<String, List<Sample>> entry : categoryToSamples.entrySet()) {
                int cur = entry.getValue().size();
                if (cur > max) { max = cur; maxCategory = entry.getKey(); }
            }
            return maxCategory;
        }
        Object[] rst = chooseBestTestAttribute(categoryToSamples, attrNames);
        Tree tree = new Tree(attrNames[(Integer) rst[0]]);
        String[] subA = new String[attrNames.length - 1];
        for (int i = 0, j = 0; i < attrNames.length; i++)
            if (i != (Integer) rst[0]) subA[j++] = attrNames[i];

        @SuppressWarnings("unchecked")
        Map<String, Map<String, List<Sample>>> splits = (Map<String, Map<String, List<Sample>>>) rst[2];
        for (Entry<String, Map<String, List<Sample>>> entry : splits.entrySet()) {
            Object attrValue = entry.getKey();
            Map<String, List<Sample>> split = entry.getValue();
            Object child = generateDecisionTree(split, subA);
            tree.setChild(attrValue, child);
        }
        return tree;
    }

以上代码涵盖了决策树构建的完整生命周期。你可以直接在本地 IDE 中运行 main 方法，观察控制台输出的树形结构，验证不同属性组合下的分类结果。