diff --git a/segment_tree/segment_tree.py b/segment_tree/segment_tree.py index cae721f..fedb5a4 100644 --- a/segment_tree/segment_tree.py +++ b/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)