Awesome F# - Decision Trees – Part II

In my previous post I went over the theory behind the ID3 algorithm. Now that we got all that painful math out of the way, let’s write some code! Here is an implementation of the algorithm in F#. (It is also attached to this blog post, download it via the link at the bottom.)

 open System

type Record = 
        Outlook     : string
        Temperature : string
        Humidity    : string
        Wind        : string
        PlayTennis  : bool 

    /// Given an attribute name return its value
    member this.GetAttributeValue(attrName) =
        match attrName with
        | "Outlook"     -> this.Outlook
        | "Temperature" -> this.Temperature
        | "Humidity"    -> this.Humidity
        | "Wind"        -> this.Wind
        | _ -> failwithf "Invalid attribute name '%s'" attrName

    /// Make the %o format specifier look all pretty like
    override this.ToString() =
            "{Outlook = %s, Temp = %s, Humidity = %s, Wind = %s, PlayTennis = %b}" 

type DecisionTreeNode =
    // Attribute name and value / child node list    
    | DecisionNode of string * (string * DecisionTreeNode) seq
    // Decision and corresponding evidence
    | Leaf         of bool * Record seq

// ----------------------------------------------------------------------------

/// Return the total true, total false, and total count for a set of Records
let countClassifications data = 
        (fun (t,f,c) item -> 
            match item.PlayTennis with
            | true  -> (t + 1, f, c + 1)
            | false -> (t, f + 1, c + 1))
        (0, 0, 0)

// ----------------------------------------------------------------------------

/// Return the theoretical number of bits required to classify the information.
/// If a 50/50 mix, returns 1, if 100% true or false returns 0.
let entropy data = 
    let (trueValues, falseValues, totalCount) = countClassifications data        

    let probTrue  = (float trueValues)  / (float totalCount)
    let probFalse = (float falseValues) / (float totalCount)

    // Log2(1.0) = infinity, short circuiting this part
    if trueValues = totalCount || falseValues = totalCount then
        -probTrue * Math.Log(probTrue, 2.0) + -probFalse * Math.Log(probFalse, 2.0)

/// Given a set of data, how many bits do you save if you know the provided attribute.
let informationGain (data : Record seq) attr =
    // Partition the data into new sets based on each unique value of the given attribute
    // e.g. [ where Outlook = rainy ], [ where Outlook = overcast], [ ... ]
    let divisionsByAttribute = 
        |> Seq.groupBy(fun item -> item.GetAttributeValue(attr))

    let totalEntropy = entropy data
    let entropyBasedOnSplit =
        |> (attributeValue, rowsWithThatValue) -> 
                        let ent = entropy rowsWithThatValue
                        let percentageOfTotalRows = (float <| Seq.length rowsWithThatValue) / (float <| Seq.length data)
                        -1.0 * percentageOfTotalRows * ent)
        |> Seq.sum

    totalEntropy + entropyBasedOnSplit
// ----------------------------------------------------------------------------

/// Give a list of attributes left to branch on and training data,
/// construct a decision tree node.
let rec createTreeNode data attributesLeft =
    let (totalTrue, totalFalse, totalCount) = countClassifications data

    // If we have tested all attributes, then label this node with the 
    // most often occuring instance; likewise if everything has the same value.
    if List.length attributesLeft = 0 || totalTrue = 0 || totalFalse = 0 then
        let mostOftenOccuring = 
            if totalTrue > totalFalse then true
            else                           false
        Leaf(mostOftenOccuring, data)
    // Otherwise, create a proper decision tree node and branch accordingly
        let attributeWithMostInformationGain =
            |> attrName -> attrName, (informationGain data attrName))
            |> List.maxBy(fun (attrName, infoGain) -> infoGain)
            |> fst
        let remainingAttributes =
            attributesLeft |> List.filter ((<>) attributeWithMostInformationGain)

        // Partition that data base on the attribute's values
        let partitionedData = 
                (fun (r : Record) -> r.GetAttributeValue(attributeWithMostInformationGain))

        // Create child nodes
        let childNodes =
            |> (fun (attrValue, subData) -> attrValue, (createTreeNode subData remainingAttributes))

        DecisionNode(attributeWithMostInformationGain, childNodes)

The entropy and informationGain functions were covered in my last post, so let’s walk through how the actual decision tree gets constructed. There’s a little work to calculating the optimal decision tree split, but with F# you can express it quite beautifully.

 let attributeWithMostInformationGain =
    |> attrName -> attrName, (informationGain data attrName))
    |> List.maxBy(fun (attrName, infoGain) -> infoGain)
    |> fst

First, it takes all the potential attributes left to split on…


… and then maps that attribute name to a new attribute name / information gain tuple …

 |> attrName -> attrName, (informationGain data attrName))

… then from the newly generated list, pick out the tuple with the highest information gain …

 |> List.maxBy(fun (attrName, infoGain) -> infoGain)

…finally returning the first element of that tuple, which is the attribute with the highest information gain.

 |> fst

Once you can construct a decision tree in memory, how do get it out? The simplest way is to print it to the console.


The code is very straight forward. Note the use of ‘padding parameter’, so that recursive calls get indented more and more. This is a very helpful technique when printing tree-like data structures to the console.

 /// Print the decision tree to the console
let rec printID3Result indent node =
    let padding = new System.String(' ', indent)

    match node with
    | Leaf(classification, data) ->
        printfn "\tClassification = %b" classification
        // data |> Seq.iter (fun item -> printfn "%s->%s" padding <| item.ToString())

    | DecisionNode(attribute, childNodes) ->
        printfn "" // Finish previous line
        printfn "%sBranching on attribute [%s]" padding attribute
        |> Seq.iter (fun (attrValue, childNode) ->
                        printf "%s->With value [%s]..." padding attrValue
                        printID3Result (indent + 4) childNode)

However, it’s almost the year 2010. So in lieu of flying cars perhaps we can at least do better than printing data to the console. Ideally, we want to generate some sexy image like this:


You could painstakingly construct the decision tree using Microsoft Visio but fortunately there are tools to do this for you. AT&T Research has produced a great tool called GraphViz. While the end result doesn’t quite have sizzle, it’s very easy enough to get going.

The following function dumps the decision tree into a format that GraphViz can plot. (Just copy the printed text into the tool and plot it using the default settings.)

 /// Prints the tree in a format amenable to GraphViz
/// See for more format
let printInGraphVizFormat node =

    let rec printNode parentName name node = 
        match node with
        | DecisionNode(attribute, childNodes) ->

            // Print the decision node
            printfn "\"%s\" [ label = \"%s\" ];" (parentName + name) attribute

            // Print link from parent to this node (unless it's the root)
            if parentName <> "" then
                printfn "\"%s\" -> \"%s\" [ label = \"%s\" ];" parentName (parentName + name) name

            |> Seq.iter(fun (attrValue, childNode) -> 
                    printNode (parentName + name) attrValue childNode)

        | Leaf(classification, _) ->
            let label =
                match classification with
                | true  -> "Yes"
                | false -> "No"
            // Print the decision node
            printfn "\"%s\" [ label = \"%s\" ];" (parentName + name) label

            // Print link from parent to this node
            printfn "\"%s\" -> \"%s\" [ label = \"%s\" ];" parentName (parentName + name) name

    printfn "digraph g {"
    printNode "" "root" node
    printfn "}"

So there you have it, ID3 in F#. With a little bit of mathematics and some clever output you can construct decision trees for all your machine learning needs. Think of the ID3 algorithm in the future the next time you want to mine customer transactions, analyze server logs, or program your killer robot to find Sarah Conner.

<TotallyShamelessPlug> If you would like to learn more about F#, check out Programming F# by O’Reilly. Available on Amazon and at other fine retailers. </TotallyShamelessPlug>