repo init - add-two-numbers complete

2025-12-06 05:17:09 -06:00 · 2021-01-17 13:18:45 -06:00
commit e2810f1ae2
1 changed files with 113 additions and 0 deletions
--- a/add-two-numbers/main.py
+++ b/add-two-numbers/main.py
@@ -0,0 +1,113 @@
+from itertools import chain
+from typing import Tuple, List, Optional
+
+import manim
+from manim import Scene, LEFT, BLUE, Square, RIGHT, VGroup, Text, DOWN, UP, Transform
+
+
+class AddTwoNumbers(Scene):
+    def construct(self):
+        squares = []
+        arrows = []
+
+        nodes1 = 8640373282
+        nodes2 = 1389430953
+
+        nodes1 = list(map(int, list(reversed(str(nodes1)))))
+        nodes2 = list(map(int, list(reversed(str(nodes2)))))
+        group1, texts1 = self.create_list_node(nodes1)
+        group2, texts2 = self.create_list_node(nodes2)
+        group1.shift(UP * 2.5)
+        group1.shift(RIGHT * 4)
+        # group2.shift(UP * 0.5)
+        group2.shift(RIGHT * 4)
+        x, y = 0, 0
+        carry = 0
+
+        final = VGroup()
+        previous_square: Optional[Square] = None
+
+        # xy_text = Text(f'X: {x} Y: {y} Carry: {carry}', size=0).to_edge(UP, buff=0.1)
+        xy_text = Text(f'Carry: {carry}', size=0.5).to_edge(UP, buff=0.1)
+        self.add(xy_text)
+        while x < len(nodes1) + 1 or y < len(nodes2) + 1:
+            # new_xy_text = Text(f'X: {x} Y: {y} Carry: {carry}', size=0.5).to_edge(UP, buff=0.1)
+            new_xy_text = Text(f'Carry: {carry}', size=0.5).to_edge(UP, buff=0.1)
+            self.play(Transform(xy_text, new_xy_text), run_time=0.2)
+
+            node_sum = carry
+            use_carry = x == len(nodes1) and y == len(nodes2) and node_sum > 0
+            if not use_carry:
+                if x < len(nodes1):
+                    node_sum += nodes1[x]
+                if y < len(nodes2):
+                    node_sum += nodes2[y]
+            carry = node_sum // 10
+
+            animations = []
+            if x < len(nodes1) + 1:
+                animations.append(group1.animate.shift(LEFT * 4))
+                x += 1
+
+            if y < len(nodes2) + 1:
+                animations.append(group2.animate.shift(LEFT * 4))
+                y += 1
+
+            if x > 0 or y > 0:
+                animations.append(final.animate.shift(LEFT * 4))
+
+            self.play(*animations)
+
+            if not use_carry:
+                if x > len(nodes1) or y > len(nodes2):
+                    continue
+
+            square = Square(stroke_color=BLUE, stroke_width=8, fill_opacity=0, stroke_opacity=1, side_length=2)
+            text = Text(str(node_sum if use_carry else node_sum % 10), size=2.6)
+
+            animations = [manim.FadeIn(square), manim.FadeIn(text)]
+            if previous_square is not None:
+                arrow = manim.Arrow(previous_square.get_right(), previous_square.get_right() + (RIGHT * 2), buff=0.05)
+                # arrow.shift(DOWN * 2.5)
+                animations.append(manim.FadeIn(arrow))
+                final.add(arrow)
+
+            text.shift(DOWN * 2.5)
+            square.shift(DOWN * 2.5)
+            self.play(*animations, run_time=0.5)
+            final.add(square)
+            final.add(text)
+            previous_square = square
+
+            self.wait(0.25)
+
+
+        self.add(group1, group2)
+        self.wait(2.5)
+
+    def create_list_node(self, nodes: List[int]) -> Tuple[VGroup, List[Text]]:
+        texts = []
+        squares = []
+        arrows = []
+        vgroup = VGroup()
+
+        for i, node in enumerate(nodes):
+            square = Square(stroke_color=BLUE, stroke_width=8, fill_opacity=0, stroke_opacity=1, side_length=2)
+            text = Text(str(node), size=2.6)
+            arrow = manim.Arrow(square.get_right(), square.get_right() + (RIGHT * 2), buff=0.05)
+
+            square.shift(RIGHT * i * 4)
+            text.shift(RIGHT * i * 4)
+            arrow.shift(RIGHT * i * 4)
+
+            texts.append(text)
+            squares.append(square)
+            if i != len(nodes) - 1:
+                arrows.append(arrow)
+
+        for object in chain.from_iterable([squares, texts, arrows]):
+            vgroup.add(object)
+        for text in texts:
+            self.add(text)
+
+        return vgroup, texts