2024-01-15 13:20:49 +01:00

119 lines
3.5 KiB
Python

#!/usr/bin/python3
from random import randint, random
import networkx as nx
import sys
DEBUG = False
class NodeInformation:
def __init__(self, color_amt):
self.available_colors = [i for i in range(0, color_amt)]
self.neighbor_colors = []
self.is_done = False
self.assigned_color = None
color_amt = None
ni = {}
g = None
def run_tick():
global color_amt, ni, g
if DEBUG: print("===== RUN TICK =====")
# run for each node
for n in g.nodes:
if ni[n].is_done:
continue
neighbors = list(g.adj[n])
# each uncolored node selects random candidate color
available_colors = ni[n].available_colors
ni[n].assigned_color = ni[n].available_colors[randint(0, len(available_colors)-1)]
if DEBUG: print(f"{n} | assigned_color: {ni[n].assigned_color}")
if DEBUG: print(f"{n} | neighbors: {neighbors}")
# nodes exchange candidate colors
for neigh in neighbors:
if not ni[neigh].is_done:
ni[neigh].neighbor_colors.append(ni[n].assigned_color)
# run again, but only after neighbor_colors assigned for all nodes
for n in g.nodes:
if ni[n].is_done:
continue
neighbors = list(g.adj[n])
if DEBUG: print(f"{n} | neighbor_colors: {ni[n].neighbor_colors}, assigned_color: {ni[n].assigned_color}")
# check if assigned color is used by neighbor
if ni[n].assigned_color in ni[n].neighbor_colors:
# reset
ni[n].assigned_color = None
ni[n].neighbor_colors = []
else:
# get permantly colored
ni[n].is_done = True
# remove color from available colors of neighbor
for neigh in neighbors:
if not ni[neigh].is_done:
if ni[n].assigned_color in ni[neigh].available_colors:
ni[neigh].available_colors.remove(ni[n].assigned_color)
is_done = True
for n in g.nodes:
is_done &= ni[n].is_done
return 0 if is_done else 1
def check_correctness():
global color_amt, ni, g
total_colors = []
for n in g.nodes:
if ni[n].assigned_color not in total_colors:
total_colors.append(ni[n].assigned_color)
neighbors = list(g.adj[n])
for neigh in neighbors:
if ni[n].assigned_color == ni[neigh].assigned_color:
print(f"FAIL: {n} and {neigh} have same color {ni[n].assigned_color}")
assert False
if len(total_colors) > color_amt:
print(f"FAIL: color_amt not correct: {len(total_colors)} <= {color_amt}")
assert False
def main():
global color_amt, ni, g
amt_nodes = randint(100, 300)
if len(sys.argv) == 2:
amt_nodes = int(sys.argv[1])
g = nx.erdos_renyi_graph(amt_nodes, random())
# find out max degree
max_degree = 0
for n in g.nodes:
max_degree = max(max_degree, g.degree[n])
color_amt = max_degree + 1
print(f"Amount of Colors: {color_amt}")
for n in g.nodes:
ni[n] = NodeInformation(color_amt)
ticks = 0
while True:
ret = run_tick()
ticks += 1
if ret == 0:
# print info
print("Final colors:")
for n in g.nodes:
print(f"{n}[{ni[n].assigned_color}] ", end="")
print()
print(f"Amount of Ticks: {ticks}")
# assert if correct
check_correctness()
exit(0)
if __name__ == '__main__':
main()