adding brace and new colors + refactoring

This commit is contained in:
Bibin Muttappillil 2022-10-14 21:17:08 +02:00
parent 39fda02b7f
commit b78789a872
1 changed files with 14 additions and 13 deletions

View File

@ -5,35 +5,36 @@ import networkx as nx
class Tree(Scene): class Tree(Scene):
def construct(self): def construct(self):
n = 2**6 n = 2**6
l, r = 42, 59
#n, l, r = 4, 1, 3 # faster test
vertices = list(range(1, 2*n)) vertices = list(range(1, 2*n))
edges = [(v // 2, v) for v in vertices[1:]] edges = [(v // 2, v) for v in vertices[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=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) graph_vis = Graph(vertices, edges, labels=False, layout="tree", layout_config={"vertex_spacing": (.2, 1.1)}, root_vertex=1).set_color("#999999")
self.play(Create(graph_vis)) self.play(Create(graph_vis), Create(Brace(VGroup(graph_vis[n+l], graph_vis[n+r-1]))))
def color_edge(v, col):
anims = [ApplyMethod(graph_vis.edges[(v // 2, v)].set_color, col)] if v != 1 else []
anims.append(ApplyMethod(graph_vis[v].set_color, col))
self.play(*anims)
def query(v, tl, tr, ql, qr): def query(v, tl, tr, ql, qr):
if tr <= ql or qr <= tl: if tr <= ql or qr <= tl:
if v != 1: color_edge(v, RED)
graph_vis.edges[(v // 2, v)].color = RED
self.play(Indicate(graph_vis[v], color=RED))
return return
if ql <= tl and tr <= qr: if ql <= tl and tr <= qr:
if v != 1: color_edge(v, GREEN)
graph_vis.edges[(v // 2, v)].color = BLUE
self.play(Indicate(graph_vis[v], color=BLUE))
return return
if v!= 1: color_edge(v, YELLOW)
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 m = (tl + tr) // 2
query(2*v, tl, m, ql, qr) query(2*v, tl, m, ql, qr)
query(2*v+1, m, tr, ql, qr) query(2*v+1, m, tr, ql, qr)
query(1, 0, n, 42, 59) query(1, 0, n, l, r)
self.wait(2) self.wait(2)