about summary refs log tree commit diff
path: root/scratch/facebook/graph-coloring.py
blob: e5b6d9c8933288c34e8f4135aa40afe5cdaaf881 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
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)))
            for c in node.neighbors:
                if c.label not in seen:
                    xs.append(c)
                    seen.add(node.label)
        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
        for c in x.neighbors:
            if c.label not in seen:
                palette.advance()
                xs.append((c, palette.get()))
                seen.add(x.label)

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)