segment tree query runtime visualization

segment_tree/segment_tree.py

@@ -4,10 +4,36 @@ import networkx as nx
 
 class Tree(Scene):
     def construct(self):
-        vertices = list(range(1, 25))
-        edges = []
-        for v in vertices[1:]:
-            edges.append((v // 2, v))
+        n = 2**6
+        vertices = list(range(1, 2*n))
+        edges = [(v // 2, v) for v in vertices[1:]]
 
-        self.play(Create(Graph(vertices, edges, labels=True, layout="tree", layout_config={"vertex_spacing": (1.0, 1.0)}, root_vertex=1)))
+        #graph_vis = Graph(vertices, edges, labels=True, layout="tree", layout_config={"vertex_spacing": (1.0, 1.0)}, root_vertex=1)
+        graph_vis = Graph(vertices, edges, labels=False, layout="tree", layout_config={"vertex_spacing": (.2, 1.1)}, root_vertex=1)
+
+        self.play(Create(graph_vis))
+
+        def query(v, tl, tr, ql, qr):
+            if tr <= ql or qr <= tl:
+                if v != 1:
+                    graph_vis.edges[(v // 2, v)].color = RED
+                self.play(Indicate(graph_vis[v], color=RED))
+                return
+            if ql <= tl and tr <= qr:
+                if v != 1:
+                    graph_vis.edges[(v // 2, v)].color = BLUE
+                self.play(Indicate(graph_vis[v], color=BLUE))
+                return
+
+            if v!= 1:
+                graph_vis.edges[(v // 2, v)].color = PURPLE
+            graph_vis[v].color = PURPLE
+            self.play(Indicate(graph_vis[v], color=YELLOW))
+
+            m = (tl + tr) // 2
+            query(2*v, tl, m, ql, qr)
+            query(2*v+1, m, tr, ql, qr)
+
+
+        query(1, 0, n, 42, 59)
         self.wait(2)