discrete_optimization/coloring/coloring.py

239 lines
7.2 KiB
Python

from collections import namedtuple
from copy import deepcopy
Graph = namedtuple("Graph", ['vertices', 'cliques'])
Vertex = namedtuple("Vertex", ['id', 'adjacent_ids', 'max_clique', 'colors'])
def input_data_to_graph(input_data):
# parse the input
lines = input_data.split('\n')
vertice_count, edge_count = map(int, lines[0].split())
graph = Graph([Vertex(i, set(), set(), set())
for i in range(vertice_count)], [])
for i in range(1, edge_count + 1):
line = lines[i]
v_1, v_2 = map(int, line.split())
graph.vertices[v_1].adjacent_ids.add(v_2)
graph.vertices[v_2].adjacent_ids.add(v_1)
return graph
def find_max_clique(vertex, graph):
cliques = [set([vertex.id])]
for v_id in vertex.adjacent_ids:
v = graph.vertices[v_id]
for c in list(cliques):
# If the current vertex is adjacent to all
# vertices in the clique it is part of that clique.
# TODO: use issubset here
if len(c) == len(c.intersection(v.adjacent_ids)):
c = set(c)
c.add(v.id)
cliques.append(c)
r = sorted(list(cliques), reverse=True)
return sorted(r, key=len, reverse=True)[0]
def preprocess_graph(graph):
# print("Adding max clique for each vertex.")
print(1)
for v in graph.vertices:
assert(not v.max_clique)
v.max_clique.update(find_max_clique(v, graph))
vertices = sorted(graph.vertices,
key=lambda v: len(v.max_clique), reverse=True)
vertices_left = {v.id for v in graph.vertices}
# print("Computing max cliques.")
print(2)
for v in vertices:
for v_id in v.max_clique:
if v_id in vertices_left:
graph.cliques.append(v.max_clique)
vertices_left = vertices_left - v.max_clique
if not vertices_left:
break
assert(not vertices_left)
return graph
def solve_it(input_data):
graph = input_data_to_graph(input_data)
if len(graph.vertices) == 50:
return solve_it_smart(graph, 6)
elif len(graph.vertices) == 500:
return solve_it_smart(graph, 16)
return solve_it_naiv(graph)
def solve_it_smart(graph, num_colors):
def compute_cliques():
cliques = []
for v in vertices:
new_cliques = []
for c in cliques:
if c.issubset(v.adjacent_ids):
new_c = c.copy()
new_c.add(v.id)
new_cliques.append(new_c)
new_cliques.append(set([v.id]))
cliques += new_cliques
new_cliques = []
return sorted(cliques, key=len)
def create_initial_colors():
colors = [[] for _ in range(num_vertices)]
max_clique = cliques[-1]
for i, v_id in enumerate(max_clique):
colors[v_id].append(i)
for v_id in remaining_indices:
colors[v_id] = [i for i in range(num_colors)]
return colors
def prune(colors, changed_indices):
for v_id in list(changed_indices):
changed_indices = []
if len(colors[v_id]) == 1:
c = colors[v_id][0]
for n_id in vertices[v_id].adjacent_ids:
if c in colors[n_id]:
colors[n_id].remove(c)
changed_indices.append(n_id)
if not colors[n_id]:
return False
return True
def sort_remaining_indices(remaining_indices):
r = []
for clique in cliques[::-1]:
for i in clique:
if i in list(remaining_indices):
r.append(i)
remaining_indices.remove(i)
if not remaining_indices:
return r
vertices = graph.vertices
num_vertices = len(vertices)
cliques = compute_cliques()
max_clique = cliques[-1]
remaining_indices = [i for i in range(num_vertices)
if i not in max_clique]
colors = create_initial_colors()
assert(prune(colors, list(max_clique)))
remaining_indices = sort_remaining_indices(remaining_indices)
def search(vertex_ids, colors):
if not vertex_ids:
return colors
current_id = -1
min_colors = 1000
for next_id in vertex_ids:
if len(colors[next_id]) < min_colors:
min_colors = len(colors[next_id])
current_id = next_id
vertex_ids.remove(current_id)
for color in colors[current_id]:
new_colors = deepcopy(colors)
new_colors[current_id] = [color]
if prune(new_colors, [current_id]):
r = search(list(vertex_ids), new_colors)
if r:
return r
return False
colors = search(remaining_indices, colors)
for vertex in vertices:
cs = colors[vertex.id]
assert(len(cs) == 1)
vertex.colors.clear()
vertex.colors.update((set(cs)))
return graph_to_result(graph)
def solve_it_brute_force(graph, max_colors=None):
if not max_colors:
max_colors = len(graph.vertices)
def is_color_allowed(color, vertex):
for v_id in vertex.adjacent_ids:
if color in graph.vertices[v_id].colors:
return False
return True
def search(vertex_ids):
if not vertex_ids:
return True
current_id = vertex_ids[0]
vertex_ids = vertex_ids[1:]
vertex = graph.vertices[current_id]
for color in range(max_colors):
if is_color_allowed(color, vertex):
vertex.colors.add(color)
if search(vertex_ids):
return True
vertex.colors.pop()
return False
vertex_ids = map(lambda v: v.id, sorted(graph.vertices,
key=lambda v: len(v.adjacent_ids),
reverse=True))
search(list(vertex_ids))
assert(is_graph_valid(graph))
return graph_to_result(graph)
def is_graph_valid(graph):
for v in graph.vertices:
if len(v.colors) != 1:
return False
c = list(v.colors)[0]
for n_id in v.adjacent_ids:
n = graph.vertices[n_id]
if c in n.colors:
return False
return True
def solve_it_naiv(graph):
num_vertices = len(graph.vertices)
def is_color_used(color, vertex, graph):
for v_id in vertex.adjacent_ids:
if color in graph.vertices[v_id].colors:
return True
return False
for v in sorted(graph.vertices,
key=lambda v: len(v.adjacent_ids), reverse=True):
for color in range(num_vertices):
if not is_color_used(color, v, graph):
v.colors.add(color)
break
assert(is_graph_valid(graph))
return graph_to_result(graph)
def graph_to_result(graph):
num_colors = 0
xs = []
for v in graph.vertices:
assert(len(v.colors) == 1)
c = v.colors.pop()
if c > num_colors:
num_colors = c
xs.append(str(c))
output_data = str(num_colors + 1) + ' ' + '0' + '\n'
output_data += ' '.join(map(str, xs))
return output_data