1 files changed, 60 insertions, 0 deletions
diff --git a/scratch/facebook/graph-coloring.py b/scratch/facebook/graph-coloring.py
new file mode 100644
index 000000000000..17588166c871
--- /dev/null
+++ b/scratch/facebook/graph-coloring.py
@@ -0,0 +1,60 @@
+from collections import deque
+
+class Palette(object):
+    def __init__(self, n):
+        self.i = 0
+        self.colors = list(range(n))
+
+    def get(self):
+        return self.colors[self.i]
+
+    def advance(self):
+        self.i += 1 % len(self.colors)
+
+class GraphNode(object):
+    def __init__(self, label):
+        self.label = label
+        self.neighbors = set()
+        self.color = None
+
+    def __repr__(self):
+        result = []
+        xs = deque()
+        xs.append(self)
+        seen = set()
+        while xs:
+            node = xs.popleft()
+            result.append('{} ({})'.format(node.label, str(node.color)))
+            seen.add(node.label)
+            for c in node.neighbors:
+                if c.label not in seen:
+                    xs.append(c)
+        return ', '.join(result)
+
+def color_graph(graph, d):
+    seen = set()
+    start = graph
+    xs = deque()
+    palette = Palette(d + 1)
+    xs.append((start, palette.get()))
+    while xs:
+        x, color = xs.popleft()
+        x.color = color
+        seen.add(x.label)
+        for c in x.neighbors:
+            if c.label not in seen:
+                palette.advance()
+                xs.append((c, palette.get()))
+
+a = GraphNode('a')
+b = GraphNode('b')
+c = GraphNode('c')
+
+a.neighbors.add(b)
+b.neighbors.add(a)
+b.neighbors.add(c)
+c.neighbors.add(b)
+
+print(a)
+color_graph(a, 3)
+print(a)