leetcode subtree visualization

subtree.py

import random

from graphviz import Digraph


# Definition for a binary tree node.
class TreeNode(object):
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None


def equal(t1, t2):
    # check if t2 is a subtree of t2
    if (t1 is None) & (t2 is None):
        return True
    elif (t1 is None) ^ (t2 is None):
        return False
    return all([
        t1.val == t2.val,
        equal(t1.left, t2.left),
        equal(t1.right, t2.right),
    ])


class Solution(object):
    def isSubtree(self, t, s):
        if equal(t, s):
            return t
        elif t is not None:
            m1 = self.isSubtree(t.left, s)
            m2 = self.isSubtree(t.right, s)
            return m1 or m2


class SubtreeViz(object):
    def __init__(self, t, s):
        self.t = t
        self.s = s

        self._common = set()
        self.preorder(self.s, self._common)

        self.dot = Digraph('Tree')

    def preorder(self, root, l):
        if root is not None:
            l.add(root.val)
            self.preorder(root.left, l)
            self.preorder(root.right, l)

    def create_node(self, node, color):
        self.dot.node(str(node.val),
                      str(node.val),
                      fillcolor=color,
                      style='filled')

    def create_edge(self, node1, node2, color=None, style=None):
        self.dot.edge(str(node1.val),
                      str(node2.val),
                      style=style,
                      fillcolor=color)

    def create_random_node(self):
        nodeid = str(hash(random.random()))
        self.dot.node(nodeid, '', style='dashed')
        return TreeNode(nodeid)

    def render_tree(self, root, color=None):
        if color is None and root is self.s:
            color = 'lightblue'
        self.create_node(root, color)
        if root.left is not None:
            self.create_edge(root, root.left, color)
            if root.right is None:
                node = self.create_random_node()
                self.create_edge(root, node, style='dashed')
            self.render_tree(root.left, color)
        if root.right is not None:
            if root.left is None:
                node = self.create_random_node()
                self.create_edge(root, node, style='dashed')
            self.create_edge(root, root.right, color)
            self.render_tree(root.right, color)

    def show(self):
        self.render_tree(self.t)
        self.dot.format = 'png'
        with open('/tmp/tmp.png', 'wb') as f:
            f.write(self.dot.pipe())


def main():
    t1 = TreeNode(3)
    t1.left = TreeNode(4)
    t1.right = TreeNode(5)
    t1.right.right = TreeNode(6)
    t1.left.left = TreeNode(1)
    t1.left.right = TreeNode(2)

    t2 = TreeNode(4)
    t2.left = TreeNode(1)
    t2.right = TreeNode(2)

    m = Solution().isSubtree(t1, t2)
    viz = SubtreeViz(t1, m)
    viz.show()


if __name__ == '__main__':
    main()