soi-manim/segment_tree/segment_tree.py

from manim import *

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).set_color("#999999")

        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:
                color_edge(v, RED)
                return
            if ql <= tl and tr <= qr:
                color_edge(v, GREEN)
                return

            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, l, r)
        self.wait(2)