simple graph coloring

color.py

#!/usr/bin/env python
import random
import sys

class node( object ):
  """ a vertex
  """
  def __init__( self, key ):
    """ c'tor """
    self.key = key 
    self.c = None
    self.outEdges = []
    self.inEdges = []

  def connect( self, other ):
    """ connect this verted to the other one  """
    if self not in other.inEdges and other not in self.outEdges:
      other.inEdges.append( self )
      self.outEdges.append( other )


class graph( object ):
  """ simple graph """

  nodes = []

  def __init__( self ):
    """ c'tor """
    self.nodes = []
    self.chroma = 1
    self.palette = []
   
  def __contains__( self, node ):
    """ in operator """
    return node in self.nodes

  def add( self, n ):
    """ add vertex """
    if n not in self.nodes:
      self.nodes.append( n )
 
  def colors( self ):
    """ create random palette """
    while len(self.palette) < self.chroma:
      newcol = "#%s" % hex( random.randint(0, 256*256*256 ) )[2:].zfill(6)
      if newcol not in self.palette:
        self.palette.append( newcol )

  def dot( self, fname ):
    """ create dot file """
    self.colors()

    if not fname.endswith(".dot"):
      fname = "%s.dot" % fname

    fname = open( fname, "w+" )
    fname.write( "digraph G {\n")

    for node in self.nodes:
      fname.write( "%s [color=\"%s\",bgcolor=\"%s\",label=\"%s\",shape=\"circle\",style=\"filled\"];\n" % ( node.key, 
                                                                                                            self.palette[node.c-1], 
                                                                                                            self.palette[node.c-1], 
                                                                                                            node.key ) )
    
    for node in self.nodes:
      for c in node.outEdges:
        fname.write( "%s -> %s;\n" % ( node.key, c.key ) )

    fname.write( "}\n" )
    fname.close()

def color( gr ):
  """ colorize graph """
  for n in gr.nodes:
    n.c = None
    
  def colorize( n, col ):
    """ can use color c for vertex n? """
    for item in n.inEdges + n.outEdges: 
      if item.c == col:
        return False
    return col

  # # sort by number of in and out edges 
  nodes = sorted( [ (len(n.outEdges + n.inEdges), n) for n in gr.nodes ], reverse = True )

  colors = [ 1 ]
  for i, n in nodes:
    if not n.c:
      while not n.c:
        for c in colors:
          can = colorize( n, c )
          if can:
            n.c = can
            break
        if not n.c:
          colors.append( colors[-1] +1 )
          n.c = colors[-1]
  ## set chromatic number
  gr.chroma = len(colors)
  return colors

## execution
if __name__ == "__main__":

  if len(sys.argv) != 4:
    print "usage: color.py nbNodes nbEdges dotFile"
    sys.exit(0)
  nodes, edges, fname = int(sys.argv[1]), int(sys.argv[2]), sys.argv[3] 


  gr = graph()

  nodes = [ node(i) for i in range(nodes) ]

  for i in range(edges):
    inode = random.choice( nodes )
    jnode = random.choice( nodes )
 
    inode.connect( jnode )

  for node in nodes:
    gr.add(node)

  c  = color(gr)
  gr.dot( fname )