anims.py

#from os import startfile
from random import randrange
from re import I
from unittest import skip
from manim import config as global_config

import solarized
from util import *



scene_width = 14.2
tree_scale = 3
node_radius = 0.2
arrow_width = 10


def create_arrow_between_nodes(start, end, node_radius, color, width):
    return Arrow(
        start=start + node_radius * np.subtract(end, start) / np.linalg.norm(np.subtract(end, start)),
        end=end - node_radius * np.subtract(end, start) / np.linalg.norm(np.subtract(end, start)),
        color=color,
        stroke_width=width,
        buff=0
    )


def sugar(scene, tree, n1, n2, n_sh):
    tree.rehang_subtree(
        scene,
        n1,
        n2,
        tree.vertices[n2].get_center() + n_sh * sh + H,
        2 * DOWN,
        2 * DOWN,
    )

class Blank(Scene):
    def construct(self):
        return

class Intro(Scene):
    def construct(self):

        num_img = 26
        img_positions = []
        n1 = 10
        n11 = 5
        n2 = 20

        for n, l1, l2 in [
            (n1, [-4, -2, 0, 2, 4], [-2, 0, 2]),
            (n11, [-4, -2, 2, 4], [-2, 0, 2]),
            (n2, [-6, -4, -2], [-2, 0, 2]),
        ]:
            for _ in range(n):
                while (True):
                    new_pos = np.array(random.choice(l1) * RIGHT + random.choice(l2) * UP)
                    i = max(0, len(img_positions) - 6)
                    collision = False
                    for pos in img_positions[i:]:
                        if np.array_equal(pos, new_pos):
                            collision = True
                    if collision == True:
                        continue
                    else:
                        img_positions.append(new_pos)
                        break

        dalle_images = []
        for i in range(len(img_positions)):
            dalle_images.append(
                ImageMobject("img/dalle/p{}.jpg".format((i % num_img) + 1)).scale_to_fit_width(2).move_to(
                    img_positions[i]
                )
            )
        dalle_title = Tex(r"Images generated by DALL·E 2", color=GRAY).scale(0.8).move_to(3.5 * UP + 3.5 * LEFT)
        gpt3_title = Tex(r"The language model GPT-3", color=GRAY).scale(0.8).next_to(dalle_title, RIGHT).shift(
            2 * RIGHT)
        shft = 7*DOWN
        dalle_website = Tex(r"https://openai.com/dall-e-2/", color = GRAY).scale(0.5).move_to(dalle_title.get_center()).shift(shft)
        gpt3_website = Tex(r"https://beta.openai.com/playground", color = GRAY).scale(0.5).move_to(gpt3_title.get_center()).shift(shft)
        dalle_title.shift(3.5 * RIGHT)
        dalle_website.shift(3.5*RIGHT)
        anims = []
        self.play(FadeIn(dalle_title), FadeIn(dalle_website))
        self.wait()

        for i in range(len(img_positions)):
            anims.append(
                Succession(
                    FadeIn(dalle_images[i]),
                    Wait(3),
                    FadeOut(dalle_images[i])
                )
            )
            if i == n1+n11:
                anims.append(AnimationGroup(
                    FadeIn(gpt3_title),
                    FadeIn(gpt3_website),
                    dalle_title.animate().shift(3.5 * LEFT),
                    dalle_website.animate().shift(3.5 * LEFT),
                ))

        self.play(
            AnimationGroup(
                *anims,
                lag_ratio=0.1
            )
        )
        self.wait()


class Polylog(Scene):
    def construct(self):
        authors = Tex(
            r"\textbf{Václav Rozhoň, Vojtěch Rozhoň, Václav Volhejn}",
            color=text_color,
            font_size=40,
        ).shift(
            3 * DOWN + 0 * LEFT
        )

        channel_name = Tex(r"polylog", color=GRAY)
        channel_name[0][1].set_color(ORANGE)
        channel_name.scale(4).shift(1 * UP)

        self.wait()

        logo_dalle = ImageMobject("img/D2.png").scale(0.5).set_z_index(100)
        logo_solarized = ImageMobject("img/logo-solarized.png").scale(0.032).shift(0.3 * UP + 0.05 * LEFT)
        prompts_dalle = [
            Tex(
                str,
                color=GRAY
            ).scale(0.8).next_to(logo_dalle, DOWN)
            for str in [r"DALL$\cdot$E 2 on the prompt:", r"A logo suitable for a youtube channel about computer science"]
        ]
        Group(*prompts_dalle).arrange(DOWN).next_to(logo_dalle, DOWN)

        asterisk = Tex(
            r"*As of 2022, neural nets generating artificial images still struggle with texts inside the images. Probably going to be solved soon. ",
            color=GRAY
        ).scale(0.2).next_to(Group(*prompts_dalle), DOWN).shift(0.4 * DOWN)

        self.play(
            FadeIn(logo_dalle),
            FadeIn(prompts_dalle[0]),
            FadeIn(prompts_dalle[1]),
            FadeIn(asterisk)
        )
        self.wait()

        self.add(logo_solarized)

        self.play(
            FadeOut(logo_dalle),
            FadeOut(prompts_dalle[0]),
            FadeOut(prompts_dalle[1]),
            FadeOut(asterisk),
        )

        self.play(
            logo_solarized.animate().move_to(2 * LEFT + 1 * UP + 0.55 * RIGHT),
            Write(authors),
            FadeIn(channel_name[0][0]),
            FadeIn(channel_name[0][2:])
        )
        self.wait()
        return

        self.play(
            *[FadeOut(o) for o in self.mobjects]
        )
        self.wait()


class Statement(Scene):
    def construct(self):
        self.next_section(skip_animations=False)

        # caption = Tex("Problem", color=GRAY).scale(3)
        # self.add_sound("audio/gong.wav")
        # self.play(
        #     FadeIn(caption)
        # )
        # self.wait()
        # self.play(
        #     FadeOut(caption)
        # )

        # The problem I’ll be solving is called Buds Re-hanging. In this problem, we are given a tree, so a bunch of nodes connected by edges such that no edges form a cycle. 

        statement = ImageMobject(
            "img/cf-statement.png",
        ).scale_to_fit_width(scene_width / 2.0).shift(scene_width / 4 * LEFT + 5.8 * DOWN)
        statement_caption = Tex("Buds Re-hanging (Codeforces 1566E)", color=GRAY).scale(0.8).move_to(
            statement.get_center()).next_to(statement, UP)

        self.play(
            FadeIn(statement),
            FadeIn(statement_caption)
        )
        self.wait()

        self.play(Circumscribe(statement_caption, color=RED))
        self.wait()

        example_tree = Tree(
            example_vertices,
            example_edges,
            layout="kamada_kawai",
            layout_scale=tree_scale,
            vertex_config={"radius": node_radius, "color": GRAY},  # for debugging
            labels=False,  # for debugging
            edge_config={"color": text_color},
            root=1
        ).move_to(scene_width / 4 * RIGHT)

        self.play(
            DrawBorderThenFill(example_tree)
        )
        self.wait()

        # This tree is also rooted, so each node except the root has one parent node and possibly some children. The nodes that have no children are called leaves – let’s highlight them in green.

        highlight_box = Rectangle(
            width=scene_width / 2 - 0.0,
            height=0.8,
            color=RED,
        ).next_to(statement_caption, DOWN).shift(1.7 * DOWN).set_z_index(100)

        self.play(
            FadeIn(highlight_box),
        )

        self.play(
            example_tree.animate().change_layout(rooted_position(pos_root=scene_width / 4 * RIGHT + 3 * UP)),
            run_time=1
        )
        self.wait()
        example_vertex = example_tree.vertices[2]
        self.play(
            example_vertex.animate().set_color(RED)
        )
        self.wait()

        # parent
        example_parent = example_tree.vertices[1]
        edge_parent = create_arrow_between_nodes(
            example_vertex.get_center(),
            example_parent.get_center(),
            node_radius,
            RED,
            arrow_width
        )
        self.play(
            Create(edge_parent),
        )
        self.play(
            Flash(example_parent, color=RED)
        )
        self.play(
            Uncreate(edge_parent)
        )
        self.wait()

        # children
        example_children = [
            example_tree.vertices[3],
            example_tree.vertices[4],
            example_tree.vertices[5],
        ]
        edge_children = [
            create_arrow_between_nodes(
                example_vertex.get_center(),
                child.get_center(),
                node_radius,
                RED,
                arrow_width
            ) for child in example_children
        ]
        self.play(
            *[Create(e) for e in edge_children],
        )
        self.play(
            *[Flash(child, color=RED) for child in example_children]
        )
        self.play(
            *[Uncreate(e) for e in edge_children],
            example_vertex.animate().set_color(GRAY)
        )
        self.wait()

        # leaves
        self.play(
            *[
                example_tree.vertices[v].animate().set_color(BLUE)
                for v in example_tree.get_leaves()
            ],
            *[
                Flash(example_tree.vertices[v], color=BLUE)
                for v in example_tree.get_leaves()
            ],
        )
        self.wait()

        # So far, these are standard terms. In this problem specifically, we also need the concept of a bud. A bud is a node that’s not a root, has at least one child, and all its children are leaves. In other words, buds are basically the nodes that have only leaves as children but are not leaves themselves. I highlighted all of these in blue.

        highlight_box.generate_target()
        highlight_box.target = Rectangle(
            width=highlight_box.get_width(),
            height=1.5,
            color=RED,
        ).next_to(highlight_box, DOWN, buff=DEFAULT_MOBJECT_TO_MOBJECT_BUFFER / 2.0).shift(0.1*UP)
        self.play(
            MoveToTarget(highlight_box),
        )
        self.wait()

        arrow = Arrow(
            start=ORIGIN + 0 * RIGHT,
            end=ORIGIN + LEFT,
            color=RED,
            stroke_width=10
        ).scale(3).move_to(3 * LEFT + 0.33  * DOWN)

        self.play(
            Create(arrow)
        )
        self.wait()
        line_height = 0.25
        for _ in range(2):
            self.play(
                arrow.animate().shift(line_height * DOWN)
            )
            self.wait()

        self.play(
            *[
                example_tree.vertices[v].animate().set_color(RED)
                for v in example_tree.get_buds()
            ],
            *[
                Flash(example_tree.vertices[v], color=RED)
                for v in example_tree.get_buds()
            ],
        )
        self.wait()

        self.play(
            Uncreate(arrow)
        )
        self.wait()


        # Now, we’re allowed to manipulate our tree it in the following way: we can take a bud and re-hang it and its children to another node of the tree. 
        self.next_section(skip_animations=False)

        highlight_box.generate_target()
        highlight_box.target = Rectangle(
            width=highlight_box.get_width(),
            height=1.3,
            color=RED,
        ).next_to(highlight_box, DOWN, buff=DEFAULT_MOBJECT_TO_MOBJECT_BUFFER / 2.0).shift(0.05*UP)
        self.play(
            MoveToTarget(highlight_box),
        )
        self.wait()

        H = 1 * DOWN

        sugar(self, example_tree, 5, 13, 0)
        sugar(self, example_tree, 5, 2, 1)


        self.wait()

        self.wait(2)

        # Notice that in this case, after we cut  this bud off the tree, this guy becomes a new bud, and after we put the bud back here, this guy stops being a leaf and also this is not a bud anymore.

        ar = Arrow(
            start=ORIGIN,
            end=(1 * RIGHT + 1 * DOWN) / 1.0,
            color=RED,
        ).move_to(
            example_tree.vertices[2].get_center() + (1 * LEFT + 1 * UP) / 2.0
        )
        ar2 = Arrow(
            start=ORIGIN,
            end=(1 * LEFT + 1 * DOWN) / 1.0,
            color=RED,
        ).move_to(example_tree.vertices[12].get_center() + (1 * RIGHT + 1 * UP) / 2.0)
        ar3 = ar2.copy().shift(H)

        self.play(
            AnimationGroup(
                FadeIn(ar),
                run_time = 0.2
            ),
        )

        anims = [
            Succession(
                Wait(0.3),
                AnimationGroup(
                    FadeIn(ar2),
                    FadeIn(ar3),
                    run_time = 0.2
                )
            )
        ]

        example_tree.rehang_subtree(
            self,
            5,
            13,
            example_tree.vertices[13].get_center() + 0 * sh + H,
            2 * DOWN,
            2 * DOWN,
            additional_anims = anims
        )
        self.play(
            FadeOut(ar),
            FadeOut(ar2),
            FadeOut(ar3),
        )
        self.wait()
        
        # And now the question is: You’re allowed to do these operations any number of times with any buds you choose. If you do the operations as cleverly as possible, what’s the lowest number of leaves the tree can have? For example, the number of leaves at the beginning is 7. You can see how it changes when we do the operations and the lowest we can get seems to be 5.

        self.next_section(skip_animations=False)

        highlight_box.generate_target()
        highlight_box.target = Rectangle(
            width=highlight_box.get_width(),
            height=0.55,
            color=RED,
        ).next_to(highlight_box, DOWN, buff=DEFAULT_MOBJECT_TO_MOBJECT_BUFFER / 2.0).shift(0.07*UP)
        self.play(
            MoveToTarget(highlight_box),
        )
        self.wait()

        num_leaves = Tex(r"\# leaves: ", color=BLUE).move_to(1.5 * RIGHT + 3 * DOWN)
        num_leaves_counter = Integer(7, color=BLUE).next_to(num_leaves, RIGHT)
        num_leaves_counter.add_updater(lambda x: x.set_value(Forest.get_leaves_cnt()))

        self.play(
            FadeIn(num_leaves),
            FadeIn(num_leaves_counter)
        )
        self.wait()

        sugar(self, example_tree, 2, 10, 0)
        sugar(self, example_tree, 5, 1, -2)
        sugar(self, example_tree, 2, 13, 0)
        sugar(self, example_tree, 5, 10, 0)
        sugar(self, example_tree, 2, 1, -2)
        sugar(self, example_tree, 5, 2, 1)
        sugar(self, example_tree, 5, 1, -5)
        sugar(self, example_tree, 12, 1, 5)
        sugar(self, example_tree, 12, 10, 0)
        sugar(self, example_tree, 5, 3, 0)
        sugar(self, example_tree, 5, 13, 0)
        sugar(self, example_tree, 2, 11, 0)
        sugar(self, example_tree, 5, 3, 0)

        # Because this is a coding problem, it is also important how large the input data is. You can see that the tree can have around 10^5 nodes, which means that our algorithm for computing the answer needs to have close to linear time complexity.

        self.play(
            Group(
                statement,
                statement_caption,
                highlight_box,
            ).animate().shift(6 * UP)
        )
        self.wait()

        highlight_box.generate_target()
        highlight_box.target = Rectangle(
            width=highlight_box.get_width(),
            height=0.6,
            color=RED,
        ).next_to(highlight_box, DOWN, buff=DEFAULT_MOBJECT_TO_MOBJECT_BUFFER / 2.0).shift(0.85 * DOWN)
        self.play(
            MoveToTarget(highlight_box),
        )
        self.wait()

        # There are examples in the problem statement to make this clearer. Let’s look at the first one: the tree looks like this. It has four leaves. But if we rehang this bud here and this bud there, we get a tree with two leaves, which is the correct answer.

        highlight_box.generate_target()
        highlight_box.target = Rectangle(
            width=highlight_box.get_width(),
            height=1.32,
            color=RED,
        ).next_to(highlight_box, DOWN, buff=DEFAULT_MOBJECT_TO_MOBJECT_BUFFER / 2.0).shift(3.2 * DOWN)
        self.play(
            MoveToTarget(highlight_box),
        )
        self.wait()
        self.play(
            Uncreate(example_tree),
        )

        sample_tree = Tree(
            sample_vertices,
            sample_edges,
            layout="kamada_kawai",
            layout_scale=tree_scale,
            vertex_config={"radius": node_radius, "color": GRAY},  # for debugging
            labels=False,  # for debugging
            edge_config={"color": text_color},
            root=1
        )  # .move_to(scene_width/4 * RIGHT)

        W = 2 * sh
        sample_tree.change_layout(
            {
                1: ORIGIN,
                2: H - W,
                3: H,
                4: H + W,
                5: 2 * H - 3 * W / 2,
                6: 2 * H - W / 2,
                7: 2 * H + W,
            }
        ).move_to(scene_width / 4 * RIGHT + 3 * UP)

        self.play(
            DrawBorderThenFill(sample_tree)
        )
        sample_tree.pretty_colour()
        self.wait()

        sugar(self, sample_tree, 4, 3, 0)
        sugar(self, sample_tree, 2, 7, 0)

        self.play(
            Group(
                statement,
                statement_caption,
                highlight_box,
            ).animate().shift(5 * UP)
        )
        highlight_box2 = Rectangle(
            width=highlight_box.get_width(),
            height=0.22,
            color=RED,
        ).move_to(highlight_box.get_center()).shift(5.18 * DOWN)
        self.play(
            Create(highlight_box2),
        )

        self.wait()

        sugar(self, sample_tree, 2, 4, 2)
        sugar(self, sample_tree, 2, 3, -2)
        sugar(self, sample_tree, 4, 1, 1)
        sugar(self, sample_tree, 2, 1, -1)
        sugar(self, sample_tree, 2, 3, 0)
        sugar(self, sample_tree, 4, 5, 0)
        sugar(self, sample_tree, 4, 3, -2)

        self.wait()

class Statement2(Scene):
    def construct(self):
        example_tree = Tree(
            example_vertices,
            example_edges,
            layout=rooted_position(),
            layout_scale=tree_scale,
            vertex_config={"radius": node_radius, "color": GRAY},  # for debugging
            labels=False,  # for debugging
            edge_config={"color": text_color},
            root=1
        ).move_to(scene_width/4 * RIGHT + 1*UP)


        self.play(
            DrawBorderThenFill(example_tree)
        )
        example_tree.pretty_colour()
        num_leaves = Tex(r"\# leaves: ", color=BLUE).move_to(1.5 * RIGHT + 3 * DOWN)
        num_leaves_counter = Integer(7, color=BLUE).next_to(num_leaves, RIGHT)
        num_leaves_counter.add_updater(lambda x: x.set_value(Forest.get_leaves_cnt()))
        self.play(FadeIn(num_leaves), FadeIn(num_leaves_counter))
        self.wait()

        sugar(self, example_tree, 5, 13, 0)
        sugar(self, example_tree, 2, 10, 0)
        sugar(self, example_tree, 5, 1, -2)
        sugar(self, example_tree, 2, 13, 0)
        sugar(self, example_tree, 5, 10, 0)
        sugar(self, example_tree, 2, 1, -2)
        sugar(self, example_tree, 5, 2, 1)

        self.wait()



class Solution(Scene):
    def construct(self):
        self.next_section(skip_animations=False)
        # caption = Tex("Solution", color=GRAY).scale(3)
        # self.add_sound("audio/gong.wav")

        # self.play(
        #     FadeIn(caption)
        # )
        # self.wait()
        # self.play(
        #     FadeOut(caption)
        # )

        num_leaves = Tex(r"\# leaves: ", color=BLUE).move_to(5.5 * LEFT + 2 * UP)
        num_leaves_counter = Integer(7, color=BLUE).next_to(num_leaves, RIGHT)

        example_tree = Tree(
            example_vertices,
            example_edges,
            layout=rooted_position(),
            layout_scale=tree_scale,
            vertex_config={"radius": node_radius, "color": GRAY},  # for debugging
            labels=False,  # for debugging
            edge_config={"color": text_color},
            root=1
        ).shift(3 * UP)

        self.play(
            DrawBorderThenFill(example_tree)
        )

        num_leaves_counter.add_updater(lambda x: x.set_value(Forest.get_leaves_cnt()))

        # self.play(
        #     FadeIn(example_tree),
        #     FadeIn(num_leaves),
        #     FadeIn(num_leaves_counter),
        # )
        # self.wait(2)

        # In competitive programming problems like these, it’s often useful to first play around with the problem and get a feeling for what the correct direction might be. So, I first drew an example tree that looked roughly like the example we already saw [obrázek stromu]
        # Then I spent a lot of time just playing with the tree simply to understand what the rehanging operations are doing. 

        sugar(self, example_tree, 5, 10, -1)
        sugar(self, example_tree, 2, 6, 0)
        sugar(self, example_tree, 2, 1, -2)
        sugar(self, example_tree, 5, 1, -5)
        sugar(self, example_tree, 12, 1, 5)
        sugar(self, example_tree, 12, 10, 0)
        sugar(self, example_tree, 5, 3, 0)
        sugar(self, example_tree, 5, 13, 0)
        sugar(self, example_tree, 2, 11, 0)
        sugar(self, example_tree, 5, 1, -2)

        # And after some time I realized the following thing. Let’s look for example at this bud and enclose it and its leaves in this pointy circle. And then do some random operations. You can see that the bud and its leaves always stay together, they never separate.

        rel_pos_5 = 0.2
        circle5 = CubicBezier(
            ORIGIN,
            1.4 * (2 * LEFT + 2 * DOWN),
            1.4 * (2 * RIGHT + 2 * DOWN),
            ORIGIN,
            color=BLACK
        )
        self.play(
            Flash(example_tree.vertices[5], color = RED)
        )
        self.wait()
        example_tree.add_object_to_vertex(5, self, circle5, rel_pos_5)

        sugar(self, example_tree, 12, 6, 0)
        sugar(self, example_tree, 2, 8, 0)
        sugar(self, example_tree, 12, 1, -5)
        sugar(self, example_tree, 2, 10, 0)
        sugar(self, example_tree, 5, 11, 0)
        sugar(self, example_tree, 12, 5, 2)

        # return back
        sugar(self, example_tree, 2, 1, -2)
        sugar(self, example_tree, 12, 10, 0)
        sugar(self, example_tree, 5, 2, 1)
        sugar(self, example_tree, 12, 11, 0)

        example_tree.remove_object(5, self)
        self.wait()

        # And that holds in general. I realized that I can repeatedly cut the buds from the tree in any order and rehang them below the root, like this. Now I draw my pointy circle around each bud and start doing some random bud-cutting operations. You can see how the nodes in the same circle always stay together. 

        sugar(self, example_tree, 5, 1, -5)
        sugar(self, example_tree, 12, 1, 5)

        circle_shifts = [
            0.4, 0.4, 0.4, 0.3
        ]
        circles = [
            CubicBezier(
                ORIGIN,
                1.4 * (2 * LEFT + 2 * DOWN),
                1.4 * (2 * RIGHT + 2 * DOWN),
                ORIGIN,
                color=BLACK
            ),
            CubicBezier(
                ORIGIN,
                1.2 * (2 * LEFT + 2 * DOWN),
                1.4 * (1.1 * RIGHT + 2 * DOWN),
                ORIGIN,
                color=BLACK
            ),
            CubicBezier(
                ORIGIN,
                1.4 * (2 * LEFT + 2 * DOWN),
                1.4 * (2 * RIGHT + 2 * DOWN),
                ORIGIN,
                color=BLACK
            ),
            CubicBezier(
                ORIGIN,
                1.3 * (1.5 * LEFT + 2 * DOWN),
                1.3 * (1.5 * RIGHT + 2 * DOWN),
                ORIGIN,
                color=BLACK
            ),
        ]

        example_tree.add_object_to_vertex(5, self, circles[0], circle_shifts[0])
        example_tree.add_object_to_vertex(2, self, circles[1], circle_shifts[1])
        example_tree.add_object_to_vertex(9, self, circles[2], circle_shifts[2])
        example_tree.add_object_to_vertex(12, self, circles[3], circle_shifts[3])
        

        sugar(self, example_tree, 5, 10, -1)
        sugar(self, example_tree, 2, 6, 0)
        sugar(self, example_tree, 12, 11, 0)
        sugar(self, example_tree, 2, 13, 0)
        sugar(self, example_tree, 5, 1, -5)        
        sugar(self, example_tree, 2, 1, -2)
        sugar(self, example_tree, 12, 3, 0)
        sugar(self, example_tree, 12, 1, 5)

        # sugar(self, example_tree, 5, 3, 0)
        # sugar(self, example_tree, 5, 4, 0)
        # sugar(self, example_tree, 5, 13, 0)
        # sugar(self, example_tree, 5, 1, -5)
        return 
        example_tree.remove_object(5, self)
        example_tree.remove_object(2, self)
        example_tree.remove_object(9, self)
        example_tree.remove_object(12, self)

        self.next_section(skip_animations=False)
        # Ok, so all the rehangings that we do are just shuffling the four buds in some ways. 
        

class Solution2(Scene):
    def construct(self):        
        
        starting_tree = Tree(
            example_vertices,
            example_edges,
            layout=rooted_position(),
            layout_scale=tree_scale,
            vertex_config={"radius": node_radius, "color": GRAY},  # for debugging
            labels=False,  # for debugging
            edge_config={"color": text_color},
            root=1
        ).shift(2 * UP + 4 * LEFT)

        self.play(
            DrawBorderThenFill(starting_tree)
        )
        starting_tree.pretty_colour()

        # When I understood this, I asked the following question. Let’s say that I take the four buds and stack them below the root in some other configuration, for example like this, to get a new tree. 

        example_tree = starting_tree.copy()
        self.play(example_tree.animate().shift(8*RIGHT))
        self.wait()
        
        buds = [
            example_tree.remove_subtree(self, v)
            for v in [5, 2, 12, 9]
        ]

        self.play(
            buds[0].animate().shift(-0 * sh + 1 * H),
            buds[1].animate().shift(2*H - sh),
            buds[2].animate().shift(1 * sh + 0*H),
            buds[3].animate().shift(0 * sh + 2 * H),
        )
        self.wait()

        self.play(
            buds[0].animate().shift(-2 * sh),
            buds[1].animate().shift(-2*H + sh),
            buds[2].animate().shift(-1 * sh + 0*H),
            buds[3].animate().shift(-2 * sh - 1 * H),
        )
        self.wait()


        example_tree.add_subtree(self, buds[1], 1)
        example_tree.add_subtree(self, buds[0], 3)
        example_tree.add_subtree(self, buds[3], 2)
        example_tree.add_subtree(self, buds[2], 9)
        self.wait()

        # The question is: Is there always a way to start with the original tree and use my rehanging operations to end up with the new tree? [pauza] Well, it turns out the answer is yes!

        a = Arrow(
            start = starting_tree.vertices[11].get_center(),
            end = example_tree.vertices[3].get_center(),
            color = RED,
            buff = 1,
        )
        a_text = Tex("Rehangings", color = GRAY).next_to(a, UP)
        self.play(FadeIn(a), FadeIn(a_text))
        self.wait()


        small_tree_scale = 0.3
        self.play(
            example_tree.animate().shift(2 * RIGHT ).scale(small_tree_scale),
            starting_tree.animate().shift(2 * LEFT ).scale(small_tree_scale),
            FadeOut(a),
            FadeOut(a_text),
        )
        self.wait()

        self.play(
            Circumscribe(starting_tree, color = RED)
        )
        self.wait()
        self.play(
            Circumscribe(example_tree, color = RED)
        )
        self.wait()

        middle_tree = Tree(
            example_vertices,
            example_edges_mid,
            layout=rooted_position_mid(),
            layout_scale=tree_scale,
            vertex_config={"radius": node_radius, "color": GRAY},  # for debugging
            labels=False,  # for debugging
            edge_config={"color": text_color},
            root=1,
        ).scale(small_tree_scale)
        middle_tree.shift(
            (example_tree.vertices[1].get_center() * 2.0 / 5 + starting_tree.vertices[1].get_center() * 3.0 / 5.0)
            - middle_tree.vertices[1].get_center()
        )

        # Why is that? Well, one fancy way of seeing it is to use my special tree where everything hangs below the root as an intermediary. 

        self.play(DrawBorderThenFill(middle_tree))
        middle_tree.pretty_colour()  # TODO
        self.wait()

        # I can convert my starting tree into this form, 

        t1 = starting_tree.copy()
        #self.play(
        #    t1.animate().shift(
        #        (middle_tree.vertices[1].get_center() - starting_tree.vertices[1].get_center()) / 2.0
        #    )
        #)
        arrow_shift = 2 * DOWN
        ar_buffer = 0.5

        shft_tree1 = shift_horizontal=(middle_tree.vertices[1].get_center() - starting_tree.vertices[1].get_center()) / 2.0

        a1 = Arrow(
            starting_tree.vertices[1].get_center() + arrow_shift,
            starting_tree.vertices[1].get_center() + shft_tree1 + arrow_shift,
            color=RED,
            buff=ar_buffer
        )
        t1.rehang_subtree(self, 5, 1, t1.vertices[1].get_center() + (-5 * sh + H) * small_tree_scale,
                          2 * DOWN * small_tree_scale, 2 * DOWN * small_tree_scale, shft_tree1
                          , additional_anims = [Create(a1)])

        t2 = t1.copy()
        # self.play(
        #     t2.animate().shift(
        #         (middle_tree.vertices[1].get_center() - starting_tree.vertices[1].get_center()) / 2.0
        #     )
        # )
        a2 = Arrow(
            t1.vertices[1].get_center() + arrow_shift,
            t1.vertices[1].get_center() + shft_tree1 + arrow_shift,
            color=RED,
            buff=ar_buffer
        )
        t2.rehang_subtree(self, 12, 1, t2.vertices[1].get_center() + (+5 * sh + H) * small_tree_scale,
                          2 * DOWN * small_tree_scale, 2 * DOWN * small_tree_scale, 
                          shift_horizontal=shft_tree1, additional_anims = [Create(a2)]
                          )

        self.wait()
        # but the same holds also for the final tree. 

        t3 = example_tree.copy()
        # self.play(
        #     t3.animate().shift(
        #         -(example_tree.vertices[1].get_center() - middle_tree.vertices[1].get_center()) / 3.0
        #     )
        # )¨
        shft_tree2 = -(example_tree.vertices[1].get_center() - middle_tree.vertices[1].get_center()) / 3.0
        a3 = Arrow(
            example_tree.vertices[1].get_center() + arrow_shift,
            example_tree.vertices[1].get_center() + shft_tree2 + arrow_shift,
            color=BLUE,
            buff=ar_buffer
        )
        t3.rehang_subtree(self, 5, 1, t3.vertices[1].get_center() + (-5 * sh + H) * small_tree_scale,
                          2 * DOWN * small_tree_scale, 2 * DOWN * small_tree_scale, shft_tree2
                                          , additional_anims = [Create(a3)]
        )


        t4 = t3.copy()
        # self.play(
        #     t4.animate().shift(
        #         -(example_tree.vertices[1].get_center() - middle_tree.vertices[1].get_center()) / 3.0
        #     )
        # )
        a4 = Arrow(
            t3.vertices[1].get_center() + arrow_shift,
            t3.vertices[1].get_center() + shft_tree2 + arrow_shift,
            color=BLUE,
            buff=ar_buffer
        )
        t4.rehang_subtree(self, 12, 1, t4.vertices[1].get_center() + (+5 * sh + H) * small_tree_scale,
                          2 * DOWN * small_tree_scale, 2 * DOWN * small_tree_scale,         shft_tree2, additional_anims = [Create(a4)]
        )


        t5 = t4.copy()
        # self.play(
        #     t5.animate().shift(
        #         -(example_tree.vertices[1].get_center() - middle_tree.vertices[1].get_center()) / 3.0
        #     )
        # )
        a5 = Arrow(
            t4.vertices[1].get_center() + arrow_shift,
            t4.vertices[1].get_center() + shft_tree2 + arrow_shift,
            color=BLUE,
            buff=ar_buffer
        )
        t5.rehang_subtree(self, 9, 1, t5.vertices[1].get_center() + (+2 * sh + H) * small_tree_scale,
                          2 * DOWN * small_tree_scale, 2 * DOWN * small_tree_scale,         shft_tree2, additional_anims = [Create(a5)]
        )
        self.wait()




        # If I now revert the operations that convert the final tree into the special form, I get a way of converting my starting tree to the final one. 
        # So, we proved that the rehanging operations can create any possible configuration of the buds. 

        tt = t5.copy()
        tt.rehang_subtree(self, 9, 2, tt.vertices[2].get_center() + (+2 * sh + H) * small_tree_scale,
                          2 * DOWN * small_tree_scale, 2 * DOWN * small_tree_scale, -shft_tree2, additional_anims = [Transform(a5, Arrow(a5.get_left(), a5.get_right(), buff=0, color=RED)),]
        )

        self.remove(tt)

        tt = t4.copy()
        tt.rehang_subtree(self, 12, 9, tt.vertices[9].get_center() + (+3 * sh + H) * small_tree_scale,
                          2 * DOWN * small_tree_scale, 2 * DOWN * small_tree_scale, -shft_tree2, additional_anims = [Transform(a4, Arrow(a4.get_left(), a4.get_right(), buff=0, color=RED)),]
        )

        self.remove(tt)

        tt = t3.copy()
        tt.rehang_subtree(self, 5, 3, tt.vertices[3].get_center() + (+0 * sh + H) * small_tree_scale,
                          2 * DOWN * small_tree_scale, 2 * DOWN * small_tree_scale, -shft_tree2, additional_anims = [Transform(a3, Arrow(a3.get_left(), a3.get_right(), buff=0, color=RED)),]
        )

        self.remove(tt)
        self.remove(middle_tree)
        self.wait()

        self.play(
            Circumscribe(starting_tree, color=RED)
        )
        self.wait()

        self.play(
            Circumscribe(example_tree, color=RED)
        )
        self.wait()

        # rubik part
        self.play(
            *[o.animate().shift(2.5 * UP) for o in
              [a1, a2, a3, a4, a5] + [t1, t2, t3, t4, t5] + [starting_tree, example_tree]]
        )

        self.wait()

        self.play(
            Circumscribe(starting_tree, color=RED)
        )
        self.wait()

        self.play(
            Circumscribe(example_tree, color=RED)
        )
        self.wait()
        self.play(
            Circumscribe(t5, color=RED)
        )
        self.wait()

        # just the main tree again
        self.play(
            *[FadeOut(a) for a in [a1, a2, a3, a4, a5]],
            *[FadeOut(t) for t in [t1, t2, t3, t4, t5, example_tree]]
        )
        self.play(
            starting_tree.animate().scale(1.0 / small_tree_scale).move_to(ORIGIN)
        )
        self.wait()



        # But that was a digression, let’s go back to our problem. By now, we actually understand perfectly what’s the deal with the rehanging operations. In fact, this whole rehanging business was just a cruel trick that the problem setter played on us. 
        self.next_section(skip_animations=False)

        # sugar(self, starting_tree, 5, 1, -5)
        # sugar(self, starting_tree, 12, 1, 5)
        # sugar(self, starting_tree, 5, 10, -1)
        # sugar(self, starting_tree, 2, 6, 0)
        # sugar(self, starting_tree, 2, 1, -2)
        # sugar(self, starting_tree, 5, 1, -5)
        # sugar(self, starting_tree, 12, 1, 5)
        # sugar(self, starting_tree, 12, 10, 0)

        # When we get the input tree, we can forget all those rehangings and just split the tree into the buds.

        h = H / 2

        sub5 = starting_tree.remove_subtree(self, 5)
        sub2 = starting_tree.remove_subtree(self, 2)
        sub12 = starting_tree.remove_subtree(self, 12)
        sub9 = starting_tree.remove_subtree(self, 9)

        self.play(
            sub5.animate().shift(1 * H + h + 2*sh - 2*sh), 
            sub9.animate().shift(2 * H + 3 * sh + h - 2*sh),
            sub2.animate().shift(2 * H + sh + h - 2*sh),
            sub12.animate().shift(h - 2*sh)
        )
        self.wait()

        pos2 = sub2.get_center()
        pos5 = sub5.get_center()
        pos9 = sub9.get_center()
        pos12 = sub12.get_center()

        cor2 = 2*sh
        cor5 = -4*sh
        cor9 = 2*sh
        cor12 = 0*sh

        # That’s because all the trees we can construct using rehangings are exactly the trees we can construct by stacking the buds below the root in any order. 

        for i in range(3):
            if i == 0:
                self.play(
                    sub2.animate().shift(3 * sh - h + cor2),
                    sub5.animate().shift(2 * sh - h + cor5),
                    sub9.animate().shift(-2 * H - 7 * sh - h + cor9),
                    sub12.animate().shift(-2 * H + sh - h + cor12),
                )
                starting_tree.add_subtree(self, sub9, 1)
                starting_tree.add_subtree(self, sub12, 1)
                starting_tree.add_subtree(self, sub5, 10)
                starting_tree.add_subtree(self, sub2, 13)

                self.wait()
                
                sub5 = starting_tree.remove_subtree(self, 5)
                sub2 = starting_tree.remove_subtree(self, 2)
                sub12 = starting_tree.remove_subtree(self, 12)
                sub9 = starting_tree.remove_subtree(self, 9)

            if i == 1:
                self.play(
                    sub2.animate().shift(-2 * H - 2 * sh - h + cor2),
                    sub5.animate().shift(-2 * H + 5 * sh - h + cor5),
                    sub9.animate().shift(-2 * sh - h + cor9),
                    sub12.animate().shift(-2 * H + 2 * sh - h + cor12),
                )
                starting_tree.add_subtree(self, sub2, 1)
                starting_tree.add_subtree(self, sub5, 1)
                starting_tree.add_subtree(self, sub12, 1)
                starting_tree.add_subtree(self, sub9, 13)

                self.wait()
                
                sub9 = starting_tree.remove_subtree(self, 9)
                sub2 = starting_tree.remove_subtree(self, 2)
                sub5 = starting_tree.remove_subtree(self, 5)
                sub12 = starting_tree.remove_subtree(self, 12)

            if i == 2:
                self.play(
                    sub2.animate().shift(-0 * H - 1 * sh - h + cor2),
                    sub5.animate().shift(-2 * H + 5 * sh - h + cor5),
                    sub9.animate().shift(-3 * sh - h + cor9),
                    sub12.animate().shift(-0 * H - 1 * sh - h + cor12),
                )

                starting_tree.add_subtree(self, sub5, 1)
                starting_tree.add_subtree(self, sub2, 6)
                starting_tree.add_subtree(self, sub12, 7)
                starting_tree.add_subtree(self, sub9, 8)

                self.wait()
                
                sub2 = starting_tree.remove_subtree(self, 2)
                sub12 = starting_tree.remove_subtree(self, 12)
                sub9 = starting_tree.remove_subtree(self, 9)
                sub5 = starting_tree.remove_subtree(self, 5)

            self.play(
                sub2.animate().move_to(pos2),
                sub5.animate().move_to(pos5),
                sub9.animate().move_to(pos9),
                sub12.animate().move_to(pos12),
            )
            self.wait()

        # Ok, so now it is a good time to recall what is the problem we actually want to solve. We want to minimize the number of leaves. So how can we do it? Well, whenever we stack two buds on top of each other, the bottom bud covers one leaf from the bud above it. 
        self.next_section(skip_animations=False)
        self.play(
            sub5.animate().shift(-2*H - h - 4*sh),
            sub2.animate().shift(-h - 1*sh)
        )
        self.wait()

        sub5.add_subtree(self, sub2, 8)
        self.wait(0.3)
        sub2 = sub5.remove_subtree(self, 2)
        self.wait(0.3)
        sub5.add_subtree(self, sub2, 8)
        self.wait(0.3)
        sub2 = sub5.remove_subtree(self, 2)

        # So, if you think about it, the best we can do is to just stack all the buds on top of each other. 

        self.play(starting_tree.animate().shift(-1.5*H + 5*sh))
        
        self.play(sub5.animate().shift(starting_tree.vertices[1].get_center() - sub5.vertices[5].get_center() + 1*H))
        starting_tree.add_subtree(self, sub5, 1)

        self.play(sub2.animate().shift(starting_tree.vertices[5].get_center() - sub2.vertices[2].get_center() + 2*H))
        starting_tree.add_subtree(self, sub2, 7)

        self.play(sub12.animate().shift(starting_tree.vertices[4].get_center() - sub12.vertices[12].get_center() + 1*H))
        starting_tree.add_subtree(self, sub12, 4)

        self.play(sub9.animate().shift(starting_tree.vertices[13].get_center() - sub9.vertices[9].get_center() + 0.3*DOWN + 2*sh))
        starting_tree.add_subtree(self, sub9, 13)

        # other_sols = Tex("There are also other solutions. ", color = GRAY).scale(0.35).shift(3.5*DOWN + 5.5*RIGHT)
        # self.play(FadeIn(other_sols))
        # self.wait()
        # self.play(FadeOut(other_sols))
        # self.wait()

        # Nice, this means there is an easy formula that computes the answer. We first split the tree into buds. Then we count the number of all the blue nodes [flash] and from that we subtract the number of red nodes minus one. 
        

        sub9 = starting_tree.remove_subtree(self, 9)
        sub12 = starting_tree.remove_subtree(self, 12)
        sub2 = starting_tree.remove_subtree(self, 2)
        sub5 = starting_tree.remove_subtree(self, 5)
        
        for pos in [pos2, pos5, pos9, pos12]:
            pos += 2*UP + 5*sh

        self.play(
            sub2.animate().move_to(pos2),
            sub5.animate().move_to(pos5),
            sub9.animate().move_to(pos9),
            sub12.animate().move_to(pos12),
        )
        self.wait()


        label_blue = Tex(r"\# blue nodes $= 8$", color = BLUE).move_to(
            4*LEFT + 2*DOWN
        )
        label_red = Tex(r"\# red nodes $= 4$", color = RED).next_to(label_blue, DOWN).align_to(label_blue, LEFT)
        label_ans = Tex(r"Answer  $= 8 - (4-1) = 5$", color = GRAY).next_to(label_red, DOWN).align_to(label_blue, LEFT)

        blue_nodes = [
            sub2.vertices[3],
            sub2.vertices[4],
            sub5.vertices[6],
            sub5.vertices[7],
            sub5.vertices[8],
            sub9.vertices[10],
            sub9.vertices[11],
            sub12.vertices[13],
        ]
        red_nodes = [
            sub2.vertices[2],
            sub5.vertices[5],
            sub9.vertices[9],
            sub12.vertices[12],
        ]
        self.play(
            *[
                Flash(v, color = BLUE) for v in blue_nodes
            ],
            FadeIn(label_blue)
        )
        self.wait()

        self.play(
            *[
                Flash(v, color = RED) for v in red_nodes
            ],
            FadeIn(label_red)
        )
        self.wait()

        pos8 = 7
        pos4 = pos8 + 3
        blue_num = Tex(r"$8$", color = BLUE).move_to(label_blue[0][-1].get_center())
        red_num = Tex(r"$4$", color = RED).move_to(label_red[0][-1].get_center())

        self.play(
            Succession(
                Write(label_ans[0][0:pos8]),
                blue_num.animate().move_to(label_ans[0][pos8].get_center()),
                Write(label_ans[0][pos8+1:pos4]),
                red_num.animate().move_to(label_ans[0][pos4].get_center()),
                Write(label_ans[0][pos4+1:]),                
            )
        )     

        
        # The minus one is because one red node has to hang below the root and each remaining one can cover one blue guy.   

        self.play(sub5.animate().shift(starting_tree.vertices[1].get_center() - sub5.vertices[5].get_center() + 1*H))
        self.play(Flash(sub5.vertices[5], color = RED))
        starting_tree.add_subtree(self, sub5, 1)

        self.play(sub2.animate().shift(starting_tree.vertices[5].get_center() - sub2.vertices[2].get_center() + 2*H))
        self.play(Flash(sub2.vertices[2], color = RED),Flash(starting_tree.vertices[7], color = BLUE))
        starting_tree.add_subtree(self, sub2, 7)

        self.play(sub12.animate().shift(starting_tree.vertices[4].get_center() - sub12.vertices[12].get_center() + 1*H))
        self.play(Flash(sub12.vertices[12], color = RED),Flash(starting_tree.vertices[4], color = BLUE))
        starting_tree.add_subtree(self, sub12, 4)

        self.play(sub9.animate().shift(starting_tree.vertices[13].get_center() - sub9.vertices[9].get_center() + 0.3*DOWN + 2*sh))
        self.play(Flash(sub9.vertices[9], color = RED),Flash(starting_tree.vertices[13], color = BLUE))
        starting_tree.add_subtree(self, sub9, 13)


        # So let’s quickly double check our formula on the tree from the sample input. 
        
        self.play(*[FadeOut(o) for o in self.mobjects])
        self.wait()

class Solution3(Scene):
    def construct(self):
        self.next_section(skip_animations=False)
        

        sample_tree = Tree(
            sample_vertices,
            sample_edges,
            layout="kamada_kawai",
            layout_scale=tree_scale,
            vertex_config={"radius": node_radius, "color": GRAY},  # for debugging
            labels=False,  # for debugging
            edge_config={"color": text_color},
            root = 1
        )  # .move_to(scene_width/4 * RIGHT)

        W = 2*sh
        sample_tree.change_layout(
            {
                1: ORIGIN,
                2: H - W,
                3: H,
                4: H + W,
                5: 2 * H - 3 * W / 2,
                6: 2 * H - W / 2,
                7: 2 * H + W,
            }
        ).move_to( 3 * UP)

        self.play(
            DrawBorderThenFill(sample_tree)
        )
        sample_tree.pretty_colour()
        self.wait()

        sub2 = sample_tree.remove_subtree(self, 2)
        self.play(sub2.animate().shift(1*H))
        
        sub4 = sample_tree.remove_subtree(self, 4)
        self.play(sub4.animate().shift(1*H))

        # There are 4 blue nodes and two red, so our formula says that the answer should be 4-1 = 3, because we can hang the buds like this. 
        
        blue_nodes = [
            sub2.vertices[5],
            sub2.vertices[6],
            sub4.vertices[7],
            sample_tree.vertices[3]
        ]
        red_nodes = [
            sub2.vertices[2],
            sub4.vertices[4]
        ]

        label_blue = Tex(r"\# blue nodes $= 4$", color = BLUE).move_to(
            4*LEFT + 2*DOWN
        )
        label_red = Tex(r"\# red nodes $= 2$", color = RED).next_to(label_blue, DOWN).align_to(label_blue, LEFT)
        label_ans = Tex(r"Answer  $= 4 - (2-1) = 3$", color = GRAY).next_to(label_red, DOWN).align_to(label_blue, LEFT)

        self.play(
            *[
                Flash(v, color = BLUE) for v in blue_nodes
            ],
            FadeIn(label_blue)
        )
        self.wait()
        
        self.play(
            *[
                Flash(v, color = RED) for v in red_nodes
            ],
            FadeIn(label_red)
        )
        self.wait()

        pos8 = 7
        pos4 = pos8 + 3
        blue_num = Tex(r"$4$", color = BLUE).move_to(label_blue[0][-1].get_center())
        red_num = Tex(r"$2$", color = RED).move_to(label_red[0][-1].get_center())

        self.play(
            Succession(
                Write(label_ans[0][0:pos8]),
                blue_num.animate().move_to(label_ans[0][pos8].get_center()),
                Write(label_ans[0][pos8+1:pos4]),
                red_num.animate().move_to(label_ans[0][pos4].get_center()),
                Write(label_ans[0][pos4+1:]),                
            )
        )     
        
        self.play(
            sub4.animate().shift(-1*H)
        )
        sample_tree.add_subtree(self, sub4, 1)
        
        self.play(
            sub2.animate().shift(1*H + 4*sh)
        )
        sample_tree.add_subtree(self, sub2, 7)
        self.wait()


        self.play(
            *[
                Flash(sample_tree.vertices[v], color = BLUE) for v in sample_tree.get_leaves()
            ],
            FadeIn(label_blue)
        )
        self.wait()        
        

        # Hm, but this is not correct, because in this case the root is kind of forming a bud with this blue node, so we can actually hang the buds like this. 
        
        self.play(
            Circumscribe(
                Group(sample_tree.vertices[1], sample_tree.vertices[3]),
                color = RED
            )
        )
        self.wait()

        sample_tree.rehang_subtree(
            self,
            4, 
            3,
            sample_tree.vertices[3].get_center() + 1*H,
            2*DOWN,
            2*DOWN
        )        
        
        # In this way, we fully utilize all the red nodes, so the answer is then just the number of blue - the number of red nodes. 

        #label_ans_new = Tex(r"Answer  $= 4 - (2-1) = 3$", color = GRAY).next_to(label_red, DOWN).align_to(label_blue, LEFT)

        self.play(
            FadeOut(label_ans[0][pos4-1]),
            FadeOut(label_ans[0][pos4+1:pos4+4]),
            red_num.animate().next_to(label_ans[0][pos4-3:pos4-1], RIGHT),
            Transform(
                label_ans[0][pos4+4:],
                Tex("$= 2$", color = GRAY).next_to(red_num, RIGHT, buff = 0.1),
            )
        )

        # So we actually need two formulas for the two different cases. 
        
        self.play(*[FadeOut(o) for o in self.mobjects])


        example_tree = Tree(
            example_vertices,
            example_edges,
            layout=rooted_position(),
            layout_scale=tree_scale,
            vertex_config={"radius": node_radius, "color": GRAY},  # for debugging
            labels=False,  # for debugging
            edge_config={"color": text_color},
            root=1
        ).shift(3*UP+3*LEFT)

        sample_tree = Tree(
            sample_vertices,
            sample_edges,
            layout="kamada_kawai",
            layout_scale=tree_scale,
            vertex_config={"radius": node_radius, "color": GRAY},  # for debugging
            labels=False,  # for debugging
            edge_config={"color": text_color},
            root = 1
        )  # .move_to(scene_width/4 * RIGHT)

        W = 2*sh
        sample_tree.change_layout(
            {
                1: ORIGIN,
                2: H - W,
                3: H,
                4: H + W,
                5: 2 * H - 3 * W / 2,
                6: 2 * H - W / 2,
                7: 2 * H + W,
            }
        ).move_to( 3 * UP + 3.5*RIGHT)

        label_ans = Tex(r"Answer = \#blue - \#red + 1", color = GRAY).next_to(example_tree, DOWN).shift(0.5*LEFT+ 1*H)
        label_ans2 = Tex(r"Answer = \#blue - \#red", color = GRAY).next_to(sample_tree, DOWN).align_to(label_ans, DOWN).shift(0.0*RIGHT)

        label_ans[0][7:12].set_color(BLUE)
        label_ans2[0][7:12].set_color(BLUE)
        label_ans[0][13:17].set_color(RED)
        label_ans2[0][13:17].set_color(RED)
        

        self.play(
            FadeIn(label_ans),
            FadeIn(label_ans2),
            DrawBorderThenFill(example_tree),
            DrawBorderThenFill(sample_tree)
        )
        self.wait()

        # decompose the right tree


        sub5 = example_tree.remove_subtree(self, 5)
        sub2 = example_tree.remove_subtree(self, 2)
        sub12 = example_tree.remove_subtree(self, 12)
        sub9 = example_tree.remove_subtree(self, 9)
        
        h = H/2

        self.play(
            sub5.animate().shift(1*H + 2*sh + h - 1.0*H - 2*sh),
            sub9.animate().shift(2*H + 3*sh + h - 1.0*H - 2*sh),
            sub2.animate().shift(2*H + sh + h - 1.0*H - 2*sh),
            sub12.animate().shift(h - 1.0*H - 2*sh)
        )
        self.wait()

        self.play(Circumscribe(example_tree.vertices[1], color = RED))
        self.wait()

        # decompose the right tree
        subr2 = sample_tree.remove_subtree(self, 2)
        self.play(subr2.animate().shift(1*H + h))
        
        subr4 = sample_tree.remove_subtree(self, 4)
        self.play(subr4.animate().shift(1*H + h))

        self.play(
            Circumscribe(Group(sample_tree.vertices[1], sample_tree.vertices[3]), color = RED)
        )
        self.wait()

        self.play(
            Circumscribe(label_ans, color = RED)
        )
        self.wait()


        self.play(
            Circumscribe(label_ans2, color = RED)
        )
        self.wait()


class AlphaCode(Scene):
    def construct(self):
        code_left = ImageMobject("img/code1.png").scale_to_fit_height(8).align_to(Dot(radius=0).move_to(7.1 * LEFT), LEFT)
        code_right = ImageMobject("img/code2.png").scale_to_fit_height(8).next_to(code_left, RIGHT, buff = 0).shift(2*LEFT)

        self.add(
            code_left,
            code_right
        )

        self.wait()

        line_height = 0.3
        rec1 = Rectangle(color = RED, height = line_height, width = 1.5).move_to(0.6*RIGHT + 0.05*DOWN)
        rec2 = Rectangle(color = RED, height = 3.6*  line_height, width = 7).move_to(0.6*RIGHT + 0.9*DOWN).align_to(rec1, LEFT)
        rec3 = Rectangle(color = RED, height = 6.7 * line_height, width = 6).move_to(4*LEFT + 0.58*UP)
        
        self.play(
            FadeIn(rec1),
            FadeIn(rec2),
            FadeIn(rec3)
        )
        self.wait()

class Dogs(Scene):
    def construct(self):
        random.seed(8)
        dogs = [
            ImageMobject("img/dogs/p{}.jpg".format(i)).scale_to_fit_width(2).move_to(
                2*UP
            ) for i in range(1, 7)
        ]
        muffins = [
            ImageMobject("img/muffins/p{}.jpg".format(i)).scale_to_fit_width(2).move_to(
                2*UP
            ) for i in range(1, 9)
        ]

        dogs_position = 2*DOWN + 3*LEFT
        for dog in dogs:
            dog.generate_target()
            dog.target.move_to(dogs_position)
        muffin_position = 2*DOWN + 3*RIGHT
        for muffin in muffins:
            muffin.generate_target()
            muffin.target.move_to(muffin_position)
                   
        dog_title = Tex("Dog", color = GRAY).next_to(dogs[0].target, DOWN)
        muffin_title = Tex("Muffin", color = GRAY).next_to(muffins[0].target, DOWN)

        self.play(
            FadeIn(dog_title),
            FadeIn(muffin_title)
        )
        self.wait()

        pics = dogs + muffins
        random.shuffle(pics)

        for i in range(len(pics)):
            pics[i].set_z_index(-i)
            pics[i].target.set_z_index(i)
            self.add(pics[i])
    
        for pic in pics:
            pic.set_z_index(i)
            self.play(MoveToTarget(pic), run_time = 0.3)
        
        self.wait()
    

class Code(Scene):
    def construct(self):
        # Nice! So we have a solution. We simply code a program that first colors the nodes blue and red so that the colors correspond to cutting the buds from our tree one by one. Then we return the number of blue nodes - the number of red nodes as the answer.        

        code = ImageMobject("img/code.png").scale_to_fit_height(8).align_to(Dot(radius=0).move_to(7.1 * LEFT), LEFT)

        dfs_brace = Brace(Line(ORIGIN, 1.6 * DOWN), RIGHT, color=GRAY).move_to(0.5 * RIGHT + 2.5 * UP)
        input_brace = Brace(Line(ORIGIN, 1.7 * DOWN), RIGHT, color=GRAY).next_to(dfs_brace, DOWN)
        run_brace = Brace(Line(ORIGIN, 0.3 * DOWN), RIGHT, color=GRAY).next_to(input_brace, DOWN)
        ans_brace = Brace(Line(ORIGIN, 0.7 * DOWN), RIGHT, color=GRAY).next_to(run_brace, DOWN)
        trick_brace = Brace(Line(ORIGIN, 0.5 * DOWN), RIGHT, color=GRAY).next_to(ans_brace, DOWN)
        testcase_brace = Brace(Line(ORIGIN, 0.9 * DOWN), RIGHT, color=GRAY).next_to(trick_brace, DOWN).shift(0.3 * DOWN)

        dfs_text = Group(
            Tex(r"Give every node a blue or a red color; ", color=GRAY).scale(0.5),
            Tex(r"a node is blue if and only if all children are red. ", color=GRAY).scale(0.5)
        ).arrange(DOWN).next_to(dfs_brace, RIGHT)
        dfs_text[0].align_to(dfs_text[1], LEFT)
        input_text = Tex(r"Read the input. ", color=GRAY).scale(0.5).next_to(input_brace, RIGHT)
        run_text = Tex(r"Run the red/blue computation.  ", color=GRAY).scale(0.5).next_to(run_brace, RIGHT)
        ans_text = Tex(r"Compute \# blue - \# red, excluding root. ", color=GRAY).scale(0.5).next_to(ans_brace, RIGHT)
        trick_text = Tex(r"Add one in case the root is blue. ", color=GRAY).scale(0.5).next_to(trick_brace, RIGHT)
        testcase_text = Tex(r"Solve all testcases. ", color=GRAY).scale(0.5).next_to(testcase_brace, RIGHT)

        self.add(
            code,
            dfs_brace,
            input_brace,
            run_brace,
            ans_brace,
            trick_brace,
            testcase_brace,
            dfs_text,
            input_text,
            run_text,
            ans_text,
            trick_text,
            testcase_text,
        )

        self.wait()

        return
        self.remove(*self.mobjects)
        self.add(code)

        sample_tree = Tree(
            sample_vertices,
            sample_edges,
            layout={
                1: ORIGIN,
                2: H - 2 * sh,
                3: H,
                4: H + 2 * sh,
                5: 2 * H - 3 * 2 * sh / 2,
                6: 2 * H - 2 * sh / 2,
                7: 2 * H + 2 * sh,
            },
            layout_scale=tree_scale,
            vertex_config={"radius": node_radius, "color": GRAY},  # for debugging
            labels=False,  # for debugging
            edge_config={"color": text_color},
            root=1,
        ).move_to(scene_width / 4 * RIGHT + 2 * UP)

        self.play(
            Create(sample_tree),
        )
        self.wait()
        sample_tree.pretty_colour()
        self.wait()

        comp = Tex(r"{{\# blue nodes}}{{ $-$ }}{{ \# red nodes }}", color=GRAY).move_to(
            sample_tree.get_center() + 4.7 * DOWN)
        comp[0].set_color(BLUE)
        comp[2].set_color(RED)
        self.play(FadeIn(comp))
        self.play(
            Transform(comp[0], Tex(r"4", color=BLUE).next_to(comp[1], LEFT)),
        )
        self.play(
            Transform(comp[2], Tex(r"3", color=RED).next_to(comp[1], RIGHT)),
        )

        comp2 = Tex(r"$= 1$", color=GRAY).next_to(comp, RIGHT)
        self.play(FadeIn(comp2))
        self.wait()
        comp3 = Tex(r"Answer = 2", color=GRAY).next_to(Group(comp, comp2), DOWN)
        self.play(FadeIn(comp3))
        self.wait()

        sugar(self, sample_tree, 4, 3, 0)
        sugar(self, sample_tree, 2, 7, 0)
        self.wait()

        ar = Arrow(
            start=ORIGIN,
            end=(1 * RIGHT) / 1.0,
            color=RED,
        ).move_to(
            sample_tree.vertices[1].get_center() + (1 * LEFT) / 2.0
        )
        self.play(
            FadeIn(ar)
        )
        self.wait()
        self.play(FadeOut(ar))
        self.wait()

        ar = Arrow(
            start=ORIGIN,
            end=(1 * LEFT) / 1.0,
            color=RED,
        ).move_to(
            2 * DOWN + 0.5 * RIGHT
        )

        self.play(
            FadeIn(ar)
        )
        self.wait()
        self.play(FadeOut(ar))
        self.wait()

        self.wait(10)

# class Thanks(Scene):
#     def construct(self):
#         s = [
#             "Big thanks to",
#             "-- 3blue1brown and Manim Community,",
#             "-- Filip Hlásek, Matěj Konečný, Hanka Rozhoňová, Tom Sláma, ",
#             "See video description for links and some more related math. :)",
#         ]
#         t = [
#             Tex(ss, color = text_color) for ss in s
#         ]
#         # for i in range(3, len(t)):
#         #     t[i].scale(0.7)
#         t[0].move_to(5*LEFT + 3*UP)
#         t[1].next_to(t[0], DOWN).align_to(t[0], LEFT)
#         t[2].next_to(t[1], DOWN).align_to(t[0], LEFT)
#         t[3].scale(0.7).next_to(t[2], DOWN).align_to(t[0], LEFT)
#         t[4].scale(0.7).next_to(t[3], DOWN).align_to(t[0], LEFT)
#         t[5].scale(0.7).next_to(t[4], DOWN).align_to(t[0], LEFT)
#         t[6].scale(0.7).next_to(t[5], DOWN).align_to(t[0], LEFT)
#         t[7].move_to(t[5].get_center()[1]*UP + 2*DOWN)

#         self.play(
#             *[FadeIn(tt) for tt in t]
#         )
#         self.wait()
#         self.play(
#             *[FadeOut(o) for o in self.mobjects]
#         )
#         self.wait()

class Explore(Scene):
    def construct(self):
        example_tree = Tree(
            example_vertices,
            example_edges,
            layout=rooted_position(),
            layout_scale=tree_scale,
            vertex_config={"radius": node_radius, "color": GRAY},  # for debugging
            labels=False,  # for debugging
            edge_config={"color": text_color},
            root=1
        ).shift(3*UP)

        sub5 = example_tree.remove_subtree(self, 5)
        sub2 = example_tree.remove_subtree(self, 2)
        sub5.add_subtree(self, sub2, 8)