SOURCE

// 决策树节点
var TreeNode = function(data, featureIndex, threshold, left, right, prediction) {
  this.data = data;           // 该节点包含的训练数据
  this.featureIndex = featureIndex;  // 分裂特征索引
  this.threshold = threshold;         // 分裂阈值
  this.left = left;                 // 左子树（<= threshold）
  this.right = right;               // 右子树（> threshold）
  this.prediction = prediction;     // 叶子节点的预测值
};

// 决策树分类器
var DecisionTree = function(maxDepth, minSamplesSplit) {
  this.maxDepth = maxDepth || 10;
  this.minSamplesSplit = minSamplesSplit || 2;
  this.root = null;
};

// 计算信息熵
DecisionTree.prototype.calculateEntropy = function(labels) {
  var total = labels.length;
  if (total === 0) return 0;

  var counts = {};
  for (var i = 0; i < total; i++) {
    var label = labels[i];
    counts[label] = (counts[label] || 0) + 1;
  }

  var entropy = 0;
  for (var label in counts) {
    var probability = counts[label] / total;
    entropy -= probability * Math.log2(probability);
  }

  return entropy;
};

// 按阈值分割数据
DecisionTree.prototype.splitData = function(data, featureIndex, threshold) {
  var left = [];
  var right = [];

  for (var i = 0; i < data.length; i++) {
    if (data[i][featureIndex] <= threshold) {
      left.push(data[i]);
    } else {
      right.push(data[i]);
    }
  }

  return { left: left, right: right };
};

// 计算信息增益
DecisionTree.prototype.informationGain = function(data, featureIndex, threshold) {
  var split = this.splitData(data, featureIndex, threshold);
  var leftLabels = split.left.map(function(row) { return row[data[0].length - 1]; });
  var rightLabels = split.right.map(function(row) { return row[data[0].length - 1]; });

  if (leftLabels.length === 0 || rightLabels.length === 0) {
    return 0;
  }

  var parentEntropy = this.calculateEntropy(
    data.map(function(row) { return row[row.length - 1]; })
  );

  var leftWeight = leftLabels.length / data.length;
  var rightWeight = rightLabels.length / data.length;

  var childEntropy = leftWeight * this.calculateEntropy(leftLabels) +
                     rightWeight * this.calculateEntropy(rightLabels);

  return parentEntropy - childEntropy;
};

// 找到最佳分裂点
DecisionTree.prototype.findBestSplit = function(data) {
  var bestGain = -1;
  var bestFeature = null;
  var bestThreshold = null;

  // 遍历所有特征（排除最后一列标签）
  var numFeatures = data[0].length - 1;

  for (var f = 0; f < numFeatures; f++) {
    // 获取该特征的所有唯一值作为候选阈值
    var values = [];
    for (var i = 0; i < data.length; i++) {
      values.push(data[i][f]);
    }
    values = values.filter(function(v, idx, arr) { return arr.indexOf(v) === idx; });
    values.sort(function(a, b) { return a - b; });

    // 用相邻值的中间点作为阈值
    for (var v = 0; v < values.length - 1; v++) {
      var threshold = (values[v] + values[v + 1]) / 2;
      var gain = this.informationGain(data, f, threshold);

      if (gain > bestGain) {
        bestGain = gain;
        bestFeature = f;
        bestThreshold = threshold;
      }
    }
  }

  return {
    featureIndex: bestFeature,
    threshold: bestThreshold,
    gain: bestGain
  };
};

// 构建树
DecisionTree.prototype.buildTree = function(data, depth) {
  // 获取当前数据的标签
  var labels = data.map(function(row) { return row[row.length - 1]; });

  // 停止条件：达到最大深度、样本太少、或所有标签相同
  if (depth >= this.maxDepth ||
      data.length < this.minSamplesSplit ||
      labels.every(function(l) { return l === labels[0]; })) {

    // 返回多数表决
    var counts = {};
    for (var i = 0; i < labels.length; i++) {
      counts[labels[i]] = (counts[labels[i]] || 0) + 1;
    }

    var maxCount = 0;
    var prediction = null;
    for (var label in counts) {
      if (counts[label] > maxCount) {
        maxCount = counts[label];
        prediction = label;
      }
    }

    return new TreeNode(data, null, null, null, null, prediction);
  }

  // 找最佳分裂
  var bestSplit = this.findBestSplit(data);

  if (bestSplit.gain <= 0) {
    // 没有有效分裂，返回多数表决
    var counts = {};
    for (var i = 0; i < labels.length; i++) {
      counts[labels[i]] = (counts[labels[i]] || 0) + 1;
    }
    var maxCount = 0;
    var prediction = null;
    for (var label in counts) {
      if (counts[label] > maxCount) {
        maxCount = counts[label];
        prediction = label;
      }
    }
    return new TreeNode(data, null, null, null, null, prediction);
  }

  // 分裂数据
  var split = this.splitData(data, bestSplit.featureIndex, bestSplit.threshold);
  var leftTree = this.buildTree(split.left, depth + 1);
  var rightTree = this.buildTree(split.right, depth + 1);

  return new TreeNode(
    data,
    bestSplit.featureIndex,
    bestSplit.threshold,
    leftTree,
    rightTree,
    null
  );
};

// 训练
DecisionTree.prototype.fit = function(data) {
  this.root = this.buildTree(data, 0);
};

// 单条数据预测
DecisionTree.prototype.predictOne = function(node, sample) {
  // 如果是叶子节点，返回预测值
  if (node.prediction !== null) {
    return node.prediction;
  }

  // 根据特征值走向左子树或右子树
  if (sample[node.featureIndex] <= node.threshold) {
    return this.predictOne(node.left, sample);
  } else {
    return this.predictOne(node.right, sample);
  }
};

// 批量预测
DecisionTree.prototype.predict = function(samples) {
  var results = [];
  for (var i = 0; i < samples.length; i++) {
    results.push(this.predictOne(this.root, samples[i]));
  }
  return results;
};

// 打印树结构（调试用）
DecisionTree.prototype.printTree = function(node, indent) {
  indent = indent || '';

  if (node.prediction !== null) {
    console.log(indent + 'Leaf: ' + node.prediction);
    return;
  }

  console.log(indent + 'Feature ' + node.featureIndex + ' <= ' + node.threshold);
  console.log(indent + '  Left:');
  this.printTree(node.left, indent + '    ');
  console.log(indent + '  Right:');
  this.printTree(node.right, indent + '    ');
};

// ========== 使用示例 ==========

// 训练数据：[年龄, 收入, 学生, 信用] -> 是否购买电脑
// 年龄: 青年=0, 中年=1, 老年=2
// 收入: 低=0, 高=1
// 学生: 否=0, 是=1
// 信用: 一般=0, 好=1
// 标签: 不购买=0, 购买=1

var trainingData = [
  [0, 0, 1, 0, 0],
  [0, 0, 1, 1, 0],
  [1, 0, 1, 0, 1],
  [2, 1, 0, 0, 1],
  [2, 1, 0, 1, 1],
  [2, 1, 0, 1, 1],
  [1, 1, 0, 1, 1],
  [0, 0, 0, 0, 0],
  [0, 1, 0, 1, 1],
  [2, 0, 0, 0, 0],
  [0, 0, 0, 1, 1],
  [1, 0, 0, 1, 1],
  [1, 0, 1, 0, 1],
  [2, 1, 1, 0, 1],
];

// 创建并训练决策树
var dt = new DecisionTree(5, 2);
dt.fit(trainingData);

// 打印树结构
console.log('决策树结构：');
dt.printTree(dt.root);

// 预测新样本
var testSamples = [
  [0, 1, 0, 1],  // 青年, 高收入, 非学生, 信用好
  [0, 0, 1, 0],  // 青年, 低收入, 是学生, 信用一般
  [2, 0, 0, 0],  // 老年, 低收入, 非学生, 信用一般
];

var predictions = dt.predict(testSamples);
console.log('\n预测结果:');
for (var i = 0; i < testSamples.length; i++) {
  console.log('  样本 ' + (i + 1) + ': ' + (predictions[i] === 1 ? '购买' : '不购买'));
}