ml-hero/Model/DecisionTree.py
2018-11-18 17:04:03 +01:00

99 lines
3.8 KiB
Python

# -*-coding:utf8 -*
from Model.Criteria import Criteria
from Model.Table import Table
from anytree import Node
class DecisionTree:
"""Classe DecisionTree permettant de contruire un arbre de décision
Atributs :
@Table table # tableau de données
@String reference # le critère à déterminer, pour lequel on construit notre arbre
@List(@Node) # Liste des noeuds
"""
def __init__(self, data, reference):
self.table = Table(data)
self.nodes = list()
self.reference = reference
# On récupère l'identité du premier noeud et on le définit
first_criteria = self._criteria_next_node(self.table)
root_node = self._add_node(first_criteria.name, None)
# Puis on lance la machine pour avancer dans l'arbre
self._create_branch(self.table, first_criteria, root_node)
def _criteria_next_node(self, table):
"""Détermine le critère en prochain noeud à partir d'un tableau"""
# On établit la liste des critères
criteria_list = table.get_criteria_list(self.reference)
# initialisation des variables
maxEntropy = 1.1
nextNode = Criteria
# On détemrine l'entropie de chaque critère
for crit in criteria_list:
criteria = Criteria(crit, table.get_column(table.table[0].index(self.reference)))
# Si c'est la plus petite jusqu'alors, c'est ce critère qui sera le prochain noeud
if criteria.get_entropy() < maxEntropy:
nextNode = criteria
maxEntropy = criteria.get_entropy()
return nextNode
def _create_branch(self, current_table, criteria, node):
"""détermine les branches partant d'un critère positionné en noeud"""
for value, counts in criteria.values.items():
# on replace le noeud parent correctement et on ajoute le critère en tant que noeud
parent_node = node
parent_node = self._add_node(value, parent_node)
# Si un total est à 0, ou qu'il n'y plus que 3 colonne dans le tableau on est en fin de branhce !
# Ou aussi qu'il n'y a qu'une seule option !
if 0 in counts.values() or len(current_table.table[0]) == 3:
result = None
# On définit le résultat en fonction de la valeur la plus importante
if counts["False"] > counts["True"]:
result = "False"
else: # Par défaut (en cas d'égalité notamment), on prend true
result = "True"
#On ajoute le noeud final
self._add_node(result, parent_node)
# sinon il faut créer un nouveau tableau à partir de l'ancien
else:
new_table = Table(current_table.remove_criteria(criteria.name, value))
# Et continuer à avancer dans l'arbre
next_criteria = self._criteria_next_node(new_table)
# Si le critère suivant n'a plus qu'une entrée possible, on est également en fin de branche !
if len(next_criteria.values) < 2:
# On l'ajoute donc en noeud final, selon sa valeur la plus importante
for values in next_criteria.values.values():
result = max(values, key=values.get)
self._add_node(result, parent_node)
# Sinon, le critère est un nouveau noeud et on continue
else:
new_parent = self._add_node(next_criteria.name, parent_node)
self._create_branch(new_table, next_criteria, new_parent)
def _add_node(self, node, parent):
"""Crée un nouveau noeud dans l'arbre. Renvoie ce nouveau noeud"""
self.nodes.append(Node(node, parent=parent))
return self.nodes[-1]