adding brace and new colors + refactoring
This commit is contained in:
parent
39fda02b7f
commit
b78789a872
|
@ -5,35 +5,36 @@ import networkx as nx
|
|||
class Tree(Scene):
|
||||
def construct(self):
|
||||
n = 2**6
|
||||
l, r = 42, 59
|
||||
#n, l, r = 4, 1, 3 # faster test
|
||||
vertices = list(range(1, 2*n))
|
||||
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=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):
|
||||
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))
|
||||
color_edge(v, 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))
|
||||
color_edge(v, GREEN)
|
||||
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))
|
||||
color_edge(v, 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)
|
||||
query(1, 0, n, l, r)
|
||||
self.wait(2)
|
||||
|
|
Loading…
Reference in New Issue