Skip to content

Instantly share code, notes, and snippets.

@cosven cosven/subtree.py
Created Apr 27, 2018

Embed
What would you like to do?
leetcode subtree visualization
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.