DecisionTree.ipynb

{
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6-final"
  },
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_data():\n",
    "    datasets = [['青年', '否', '否', '一般', '否'],\n",
    "               ['青年', '否', '否', '好', '否'],\n",
    "               ['青年', '是', '否', '好', '是'],\n",
    "               ['青年', '是', '是', '一般', '是'],\n",
    "               ['青年', '否', '否', '一般', '否'],\n",
    "               ['中年', '否', '否', '一般', '否'],\n",
    "               ['中年', '否', '否', '好', '否'],\n",
    "               ['中年', '是', '是', '好', '是'],\n",
    "               ['中年', '否', '是', '非常好', '是'],\n",
    "               ['中年', '否', '是', '非常好', '是'],\n",
    "               ['老年', '否', '是', '非常好', '是'],\n",
    "               ['老年', '否', '是', '好', '是'],\n",
    "               ['老年', '是', '否', '好', '是'],\n",
    "               ['老年', '是', '否', '非常好', '是'],\n",
    "               ['老年', '否', '否', '一般', '否'],\n",
    "               ]\n",
    "    labels = [u'年龄', u'有工作', u'有自己的房子', u'信贷情况', u'类别']\n",
    "    # 返回数据集和每个维度的名称\n",
    "    return datasets, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "    年龄 有工作 有自己的房子 信贷情况 类别\n0   青年   否      否   一般  否\n1   青年   否      否    好  否\n2   青年   是      否    好  是\n3   青年   是      是   一般  是\n4   青年   否      否   一般  否\n5   中年   否      否   一般  否\n6   中年   否      否    好  否\n7   中年   是      是    好  是\n8   中年   否      是  非常好  是\n9   中年   否      是  非常好  是\n10  老年   否      是  非常好  是\n11  老年   否      是    好  是\n12  老年   是      否    好  是\n13  老年   是      否  非常好  是\n14  老年   否      否   一般  否\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "datasets, labels = create_data()\n",
    "train_data = pd.DataFrame(datasets, columns=labels)\n",
    "print(train_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "0.08300749985576883\n0.32365019815155627\n0.4199730940219749\n0.36298956253708536\n"
     ]
    }
   ],
   "source": [
    "def CalEnt(data):\n",
    "    total = len(data)\n",
    "    cate_sets = set(data['类别'])\n",
    "    ent = 0\n",
    "    for cate in cate_sets:\n",
    "        sub_data = data[data['类别']==cate]\n",
    "        ent -= len(sub_data)/total*math.log(len(sub_data)/total,2)\n",
    "    return ent\n",
    "\n",
    "def InfoGain(data,feature_name):\n",
    "    total = len(data)\n",
    "    total_ent = CalEnt(data)\n",
    "    feature_sets = set(data[feature_name])\n",
    "    ent = 0\n",
    "    for feature in feature_sets:\n",
    "        sub_data = data[data[feature_name]==feature]\n",
    "        ent += len(sub_data)/total*CalEnt(sub_data)\n",
    "    return total_ent - ent\n",
    "\n",
    "def InfoGainRatio(data,feature_name):\n",
    "    total = len(data)\n",
    "    total_ent = CalEnt(data)\n",
    "    feature_sets = set(data[feature_name])\n",
    "    ent = 0\n",
    "    den = 0\n",
    "    for feature in feature_sets:\n",
    "        sub_data = data[data[feature_name]==feature]\n",
    "        ent += len(sub_data)/total*CalEnt(sub_data)\n",
    "        den -= len(sub_data)/total*math.log(len(sub_data)/total,2)\n",
    "    return (total_ent - ent)/den\n",
    "\n",
    "\n",
    "print(InfoGain(train_data,\"年龄\"))\n",
    "print(InfoGain(train_data,\"有工作\"))\n",
    "print(InfoGain(train_data,\"有自己的房子\"))\n",
    "print(InfoGain(train_data,\"信贷情况\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TreeNode(object):\n",
    "    def __init__(self,split_feature=None,children=None,cate=None):\n",
    "        self.split_feature = split_feature\n",
    "        self.children = dict()\n",
    "        self.cate = cate\n",
    "        \n",
    "class DesitionTree(object):\n",
    "    def __init__(self,train,labels,ep,criterion=InfoGainRatio):\n",
    "        self.train = train\n",
    "        self.root = TreeNode()\n",
    "        self.labels = labels\n",
    "        self.ep = ep\n",
    "        self.criterion = criterion\n",
    "\n",
    "    def BuildTree(self,data):\n",
    "        if len(data) == 0:\n",
    "            return None\n",
    "        if set(data['类别'])==1:\n",
    "            return TreeNode(cate = data['类别'].mode()[0])\n",
    "        info_gain = []\n",
    "        for feature_name in self.labels[:-1]:\n",
    "            info_gain.append(self.criterion(data,feature_name))\n",
    "        if max(info_gain)<self.ep:\n",
    "            return TreeNode(cate=data['类别'].mode()[0])\n",
    "        split_feature = self.labels[info_gain.index(max(info_gain))]\n",
    "        node = TreeNode(split_feature = split_feature,cate=data['类别'].mode()[0])\n",
    "        for val in set(data[split_feature]):\n",
    "            node.children[val]=self.BuildTree(data[data[split_feature]==val])\n",
    "        return node\n",
    " \n",
    "    def Train(self):\n",
    "        self.root = self.BuildTree(self.train)\n",
    "\n",
    "    def Predict(self,node,sample):\n",
    "        if node.split_feature == None or not node.children:\n",
    "            return node.cate\n",
    "        return self.Predict(node.children[sample[node.split_feature]],sample)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "True\nTrue\nTrue\nTrue\nTrue\nTrue\nTrue\nTrue\nTrue\nTrue\nTrue\nTrue\nTrue\nTrue\nTrue\n"
     ]
    }
   ],
   "source": [
    "tree = ID3(train_data,labels,0.01)\n",
    "tree.Train()\n",
    "for idx, sample in train_data.iterrows():\n",
    "    print(sample['类别']==tree.Predict(tree.root,sample))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      " 有自己的房子:否 有工作:否 --> 否\n 有自己的房子:否 有工作:是 --> 是\n 有自己的房子:是 --> 是\n"
     ]
    }
   ],
   "source": [
    "def dfs(node,tmp):\n",
    "    if node!= None:\n",
    "        if node.split_feature:\n",
    "            tmp = tmp + \" \"+ node.split_feature\n",
    "        if not node.split_feature:\n",
    "            print((tmp+\" --> \"+node.cate))\n",
    "    else:\n",
    "        return\n",
    "    for k in node.children:\n",
    "        dfs(node.children[k],tmp+\":\"+k)\n",
    "dfs(tree.root,\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ]
}