1 files changed, 71 insertions, 0 deletions
diff --git a/scratch/facebook/mst.py b/scratch/facebook/mst.py
new file mode 100644
index 000000000000..81aa5cd48744
--- /dev/null
+++ b/scratch/facebook/mst.py
@@ -0,0 +1,71 @@
+from heapq import heappush, heappop
+import random
+
+def to_vertex_list(graph):
+    result = {}
+    for a, b, kg in graph:
+        if a in result:
+            result[a].append((b, kg))
+        else:
+            result[a] = [(b, kg)]
+        if b in result:
+            result[b].append((a, kg))
+        else:
+            result[b] = [(a, kg)]
+    return result
+
+def mst(graph):
+    graph = to_vertex_list(graph)
+    beg = random.choice(list(graph.keys()))
+    h = []
+    result = []
+    seen = set()
+    for c, kg in graph[beg]:
+        heappush(h, (kg, beg, c))
+    while h:
+        kg, beg, end = heappop(h)
+        # detect cycles
+        if end in seen:
+            continue
+        # use the edge
+        seen.add(beg)
+        seen.add(end)
+        result.append((beg, end))
+        for c, kg in graph[end]:
+            heappush(h, (kg, end, c))
+    return result
+
+graphs = [
+    [
+        ('A', 'B', 7),
+        ('A', 'D', 5),
+        ('B', 'D', 9),
+        ('E', 'D', 15),
+        ('F', 'D', 6),
+        ('F', 'G', 11),
+        ('F', 'E', 8),
+        ('G', 'E', 9),
+        ('C', 'E', 5),
+        ('B', 'E', 7),
+        ('B', 'C', 8),
+    ],
+    [
+        ('A', 'B', 4),
+        ('A', 'C', 8),
+        ('B', 'C', 11),
+        ('B', 'E', 8),
+        ('C', 'D', 7),
+        ('C', 'F', 1),
+        ('D', 'E', 2),
+        ('D', 'F', 6),
+        ('E', 'G', 7),
+        ('E', 'H', 4),
+        ('F', 'H', 2),
+        ('G', 'H', 14),
+        ('G', 'I', 9),
+        ('H', 'I', 10),
+    ],
+]
+
+for graph in graphs:
+    print(mst(graph))