1. import matplotlib.pyplot as plt
    3. decisionNode = dict(boxstyle="sawtooth", fc="0.8")
    4. leafNode = dict(boxstyle="round4", fc="0.8")
    5. arrow_args = dict(arrowstyle="<-")
    7. def plotNode(nodeTxt, centerPt, parentPt, nodeType):
    8.     createPlot.ax1.annotate(nodeTxt, xy=parentPt, xycoords='axes fraction', \
    9.                             xytext=centerPt, textcoords='axes fraction', \
    10.                             va="center", ha="center", bbox=nodeType, arrowprops=arrow_args)
    12. def getNumLeafs(myTree):
    13.     numLeafs = 0
    14.     firstStr = list(myTree.keys())[0]
    15.     secondDict = myTree[firstStr]
    16.     for key in secondDict.keys():
    17.         if type(secondDict[key]).__name__ == 'dict':
    18.             numLeafs += getNumLeafs(secondDict[key])
    19.         else:
    20.             numLeafs += 1
    21.     return numLeafs
    23. def getTreeDepth(myTree):
    24.     maxDepth = 0
    25.     firstStr = list(myTree.keys())[0]
    26.     secondDict = myTree[firstStr]
    27.     for key in secondDict.keys():
    28.         if type(secondDict[key]).__name__ == 'dict':
    29.             thisDepth = getTreeDepth(secondDict[key]) + 1
    30.         else:
    31.             thisDepth = 1
    32.         if thisDepth > maxDepth:
    33.             maxDepth = thisDepth
    34.     return maxDepth
    36. def plotMidText(cntrPt, parentPt, txtString):
    37.     xMid = (parentPt[0] - cntrPt[0]) / 2.0 + cntrPt[0]
    38.     yMid = (parentPt[1] - cntrPt[1]) / 2.0 + cntrPt[1]
    39.     createPlot.ax1.text(xMid, yMid, txtString)
    41. def plotTree(myTree, parentPt, nodeTxt):
    42.     numLeafs = getNumLeafs(myTree)
    43.     depth = getTreeDepth(myTree)
    44.     firstStr = list(myTree.keys())[0]
    45.     cntrPt = (plotTree.xOff + (1.0 + float(numLeafs)) / 2.0 / plotTree.totalw, plotTree.yOff)
    46.     plotMidText(cntrPt, parentPt, nodeTxt)
    47.     plotNode(firstStr, cntrPt, parentPt, decisionNode)
    48.     secondDict = myTree[firstStr]
    49.     plotTree.yOff = plotTree.yOff - 1.0 / plotTree.totalD
    50.     for key in secondDict.keys():
    51.         if type(secondDict[key]).__name__ == 'dict':
    52.             plotTree(secondDict[key], cntrPt, str(key))
    53.         else:
    54.             plotTree.xOff = plotTree.xOff + 1.0 / plotTree.totalw
    55.             plotNode(secondDict[key], (plotTree.xOff, plotTree.yOff), cntrPt, leafNode)
    56.             plotMidText((plotTree.xOff, plotTree.yOff), cntrPt, str(key))
    57.     plotTree.yOff = plotTree.yOff + 1.0 / plotTree.totalD
    59. def createPlot(inTree):
    60.     fig = plt.figure(1, facecolor='white')
    61.     fig.clf()
    62.     axprops = dict(xticks=[], yticks=[])
    63.     createPlot.ax1 = plt.subplot(111, frameon=False, **axprops)
    64.     plotTree.totalw = float(getNumLeafs(inTree))
    65.     plotTree.totalD = float(getTreeDepth(inTree))
    66.     plotTree.xOff = -0.5 / plotTree.totalw
    67.     plotTree.yOff = 1.0
    68.     plotTree(inTree, (0.5, 1.0), '')
    69.     plt.show()