Update app.py

app.py

@@ -1,8 +1,6 @@
 import gradio as gr
 import matplotlib.pyplot as plt
 import networkx as nx
-import io
-from PIL import Image
 
 # Node class for tree structures
 class Node:
@@ -22,16 +20,28 @@ def insert(root, key):
     return root
 
 # Function to perform inorder traversal
-def inorder(root):
-    return inorder(root.left) + [root.val] + inorder(root.right) if root else []
+def inorder(root, traversal=[]):
+    if root:
+        inorder(root.left, traversal)
+        traversal.append(root.val)
+        inorder(root.right, traversal)
+    return traversal
 
 # Function to perform preorder traversal
-def preorder(root):
-    return [root.val] + preorder(root.left) + preorder(root.right) if root else []
+def preorder(root, traversal=[]):
+    if root:
+        traversal.append(root.val)
+        preorder(root.left, traversal)
+        preorder(root.right, traversal)
+    return traversal
 
 # Function to perform postorder traversal
-def postorder(root):
-    return postorder(root.left) + postorder(root.right) + [root.val] if root else []
+def postorder(root, traversal=[]):
+    if root:
+        postorder(root.left, traversal)
+        postorder(root.right, traversal)
+        traversal.append(root.val)
+    return traversal
 
 # Function to perform level order traversal
 def level_order(root):
@@ -48,9 +58,10 @@ def level_order(root):
             queue.append(node.right)
     return result
 
-# Function to visualize the tree structure
-def visualize_tree(root):
+# Function to visualize the tree structure with traversal
+def visualize_tree_with_traversal(root, traversal):
     graph = nx.DiGraph()
+    labels = {}
     
     def add_edges(node, pos=None):
         if node is not None:
@@ -62,15 +73,15 @@ def visualize_tree(root):
     
     add_edges(root)
     
+    # Label the nodes with their position in the traversal
+    for idx, val in enumerate(traversal):
+        labels[val] = f'{val}\n({idx+1})'
+
     plt.figure(figsize=(10, 6))
     pos = hierarchy_pos(graph, root.val)
-    nx.draw(graph, pos, with_labels=True, node_size=5000, node_color="skyblue", font_size=16, font_weight="bold", font_color="black")
-    buf = io.BytesIO()
-    plt.savefig(buf, format='png')
-    buf.seek(0)
-    img = Image.open(buf)
-    plt.close()
-    return img
+    nx.draw(graph, pos, with_labels=True, labels=labels, node_size=5000, node_color="skyblue", font_size=16, font_weight="bold", font_color="black")
+    plt.show()
+    return plt
 
 # Helper function for hierarchical positioning
 def hierarchy_pos(G, root=None, width=1., vert_gap = 0.2, vert_loc = 0, xcenter = 0.5):
@@ -94,25 +105,24 @@ def _hierarchy_pos(G, root, width=1., vert_gap = 0.2, vert_loc = 0, xcenter = 0.
             pos = _hierarchy_pos(G, child, width=dx, vert_gap=vert_gap, vert_loc=vert_loc-vert_gap, xcenter=nextx, pos=pos, parent=root, parsed=parsed)
     return pos
 
-# Function to perform tree operations and visualization
+# Function to perform tree operations and visualization with traversal order
 def perform_tree_operations(values, operation):
     root = None
     for value in values:
         root = insert(root, value)
     
+    traversal = []
     if operation == "Inorder Traversal":
-        result = inorder(root)
+        traversal = inorder(root)
     elif operation == "Preorder Traversal":
-        result = preorder(root)
+        traversal = preorder(root)
     elif operation == "Postorder Traversal":
-        result = postorder(root)
+        traversal = postorder(root)
     elif operation == "Level Order Traversal":
-        result = level_order(root)
-    else:
-        result = None
+        traversal = level_order(root)
 
-    img = visualize_tree(root)
-    return result, img
+    plt = visualize_tree_with_traversal(root, traversal)
+    return None, plt
 
 # Gradio interface setup
 with gr.Blocks() as demo:
@@ -125,16 +135,15 @@ with gr.Blocks() as demo:
             submit = gr.Button("Submit")
         
         with gr.Column():
-            output = gr.Textbox(label="Output", lines=4)
-            graph_output = gr.Image(label="Tree Visualization")
+            graph_output = gr.Plot(label="Tree Visualization")
 
     # Function to handle Gradio interaction
     def on_submit(values, operation):
         values = list(map(int, values.split(',')))
-        result, img = perform_tree_operations(values, operation)
-        return str(result), img
+        _, plt = perform_tree_operations(values, operation)
+        return plt
 
-    submit.click(on_submit, inputs=[input_values, operation], outputs=[output, graph_output])
+    submit.click(on_submit, inputs=[input_values, operation], outputs=[graph_output])
 
 # Launch the Gradio app
 demo.launch()