Files
visuals/add-two-numbers/main.py
2021-01-17 13:18:45 -06:00

114 lines
3.8 KiB
Python

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