Search code examples
pythonpython-3.xtreedecision-treediagram

How to create and traverse decision tree in Python


I am trying to create a Tree in python that looks like this diagram:

Diagram

How do I write code in Python to create this tree and find a given "answer" or leaf node, given gender and color.

For example:

Input: { Gender : "Female", Color : "Green" }

Output: "Message 5"

I am going to add more levels and nodes, so I trying to create a tree representation rather than a bunch of "if" statements, since that could be messy.


Solution

  • You can create the tree in a dict structure, e.g.:

    tree = {'Male': {'Red': 'Message 1', 'Green': 'Message 2', 'Blue': 'Message 3'}, 
            'Female': {'Red': 'Message 4', 'Green': 'Message 5', 'Blue': 'Message 6'}}
    

    Then traversing this tree is nothing more than key lookups, e.g.:

    In []:
    i = {'Gender' : "Female", 'Color' : "Green" }
    tree[i['Gender']][i['Color']]
    
    
    Out[]:
    "Message 5"