karger.py

from collections import Counter
from manim import *
import networkx as nx
import random
import numpy as np
import itertools


class UF:
    def __init__(self, ids):
        self.name = {i: i for i in ids}

    def find(self, i):
        while i != self.name[i]:
            i = self.name[i]
        return i

    def union(self, i, j):
        # j -> i
        i = self.find(i)
        j = self.find(j)
        self.name[j] = i


def transform_vdicts(scene, vdict1, vdict2, run_time=1):
    k1 = vdict1.submob_dict.keys()
    k2 = vdict2.submob_dict.keys()
    transforms = []
    for key in k1 - k2:
        transforms.append(FadeOut(vdict1[key]))
    for key in k2 - k1:
        transforms.append(Create(vdict2[key]))
    for key in k1 & k2:
        transforms.append(Transform(vdict1[key], vdict2[key]))
    scene.play(*transforms, run_time=run_time)


class KargerMinCut(Scene):
    def construct(self):
        # Create a random graph
        S = 6
        G = nx.MultiGraph()
        N = 10
        edges = []
        for i in range(1, N):
            for j in range(1, N):
                if random.random() < 30 / N**2:
                    edges.append((i, j))

        G.add_edges_from(edges)

        def draw_graph(og, layout, color_edge, names):
            layout = {node: S * np.array([x, y, 0]) for node, (x, y) in layout.items()}
            edges = []

            total = Counter()
            for u, v, idx in og.edges:
                i, j = names.find(u), names.find(v)
                total[min(i, j), max(i, j)] += 1

            vdict = VDict()

            if color_edge is not None:
                uu, vv = color_edge
                ii, jj = names.find(uu), names.find(vv)
                color_edge = (min(ii, jj), max(ii, jj))
                col_idx = random.randrange(total[color_edge]) + 1

            cnt = Counter()
            for u, v, idx in og.edges:
                i, j = names.find(u), names.find(v)

                # Have to get the start and end of the edge before normalizing (i,j)
                # since otherwise the edge may "do a flip" in the animation
                start, end = layout[i], layout[j]

                edge = (min(i, j), max(i, j))
                cnt[edge] += 1

                # Bend the edge if it is a multi-edge
                control1, control2 = start, end
                if i != j:
                    rad = cnt[edge] // 2 * (-1) ** cnt[edge] * 0.05
                    if total[edge] % 2 == 0:
                        rad -= 0.025
                    if i > j:
                        rad *= -1
                    vec = end - start
                    vec = np.array([[0,1,0],[-1,0,0],[0,0,0]]) @ vec
                    control1 = start + vec * rad
                    control2 = end + vec * rad

                # Color the edge red
                color = WHITE
                stroke_width = 2
                if color_edge is not None:
                    if edge == color_edge and cnt[edge] == col_idx:
                        color = RED
                        stroke_width = 8

                vdict[u, v, idx] = CubicBezier(
                    start_anchor=start,
                    end_anchor=end,
                    start_handle=control1,
                    end_handle=control2,
                    # start_handle=control1*5/6 + control2/6,
                    # end_handle=control2*5/6 + control1/6,
                    stroke_width=S * stroke_width,
                    color=color,
                )

            for node in og.nodes:
                i = names.find(node)
                vdict[node] = Dot(point=layout[i], radius=S * 0.08, color=BLUE)

            return vdict

        pos = nx.shell_layout(G)
        names = UF(G.nodes)
        graph_vgroup = draw_graph(G, pos, None, names)
        self.add(graph_vgroup)
        self.play(Create(graph_vgroup))

        original_G = G.copy()

        while len(G.nodes) > 2:
            while True:
                u, v, idx = random.choice(list(G.edges))
                if u != v:
                    break

            # Color the edge red
            new_graph_vgroup = draw_graph(original_G, pos, (u, v), names)
            transform_vdicts(self, graph_vgroup, new_graph_vgroup, run_time=1)

            # Contract the edge
            G = nx.contracted_edge(G, (u, v, idx))
            new_pos = nx.shell_layout(G)
            names.union(u, v)  # v is merged into u

            new_graph_vgroup = draw_graph(original_G, new_pos, (u, v), names)
            transform_vdicts(self, graph_vgroup, new_graph_vgroup, run_time=1)
            pos = new_pos

        self.wait()