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 _ in range(30):
            edges.append((random.randint(1, N), random.randint(1, N)))
        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 = []
            cnt = Counter()

            vdict = VDict()

            for u, v, idx in og.edges:
                i, j = names.find(u), names.find(v)
                start, end = layout[i], layout[j]

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

                # Bend the edge if it is a multi-edge
                control1, control2 = start, end
                if i != j:
                    rad = cnt[i, j] // 2 * (-1) ** cnt[i, j] * 0.1 * S
                    control1 = start + np.array([0, rad, 0])
                    control2 = end + np.array([0, rad, 0])

                # Color the edge red
                color = WHITE
                stroke_width = 2
                if color_edge is not None:
                    uu, vv = color_edge
                    ii, jj = names.find(uu), names.find(vv)
                    ii, jj = min(ii, jj), max(ii, jj)
                    if (i, j) == (ii, jj) and cnt[i, j] == 1:
                        color = RED
                        stroke_width = 8

                vdict[u, v, idx] = CubicBezier(
                    start_anchor=start,
                    end_anchor=end,
                    start_handle=control1,
                    end_handle=control2,
                    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()