Count in how many ways you can label an unrooted tree (graph theory)

count_trees.py

# coding=UTF-8
from math import factorial
from collections import defaultdict

__author__ = u'Piotr Maślanka, EF-DU/AA, 127172'

class Graph(object):
    """Immutable graph object"""

    def __init__(self, links=set()):
        """:param links set of tuples (vertex1, vertex2)"""
        self.links = links

    def get_vertices_from(self, vertex, sans=None):
        """Return list of vertices reachable from vertex.
        If sans is set, this is a vertex that will not be included in result set"""
        vs = []
        for v1, v2 in self.links:
            if v1 == vertex:
                if v2 != sans:
                    vs.append(v2)
            elif v2 == vertex:
                if v1 != sans:
                    vs.append(v1)

        return vs

    def trace_tree(self, start_vertex, towards_vertex):
        """Assuming this tree is a graph, return a subtree.
        Return all vertices in this tree that are connected to towards_vertex, but don't check vertices
        that go toward start_vertex.

        start_vertex will not be included, but towards_vertex will
        """
        cg = Graph()

        def track(g, sv, tv):
            for next_vertex in g.get_vertices_from(tv, sans=sv):
                yield tv, next_vertex
                for s in track(g, tv, next_vertex):
                    yield s

        pg = set(track(self, start_vertex, towards_vertex))
        if len(pg) == 0:
            return SingleVertex()
        else:
            return Graph(pg)

    def split_tree(self, vertex):
        """Return rooted trees that were caused by splitting this tree into subtrees
        by removing vertex.
        Amount of trees returned will be equal to degree of vertex.
        :return a list of tuple (root vertex, subtree)"""
        t = []
        for next_vertex in self.get_vertices_from(vertex):
            t += [(next_vertex, self.trace_tree(vertex, next_vertex))]
        return t

    def get_degree(self, vertex):
        """Return degree of a vertex"""
        return len(self.get_vertices_from(vertex))

    def get_vertices(self):
        """Return a set of all vertices"""
        vs = set()
        for v1, v2 in self.links:
            vs.add(v1)
            vs.add(v2)
        return vs

    def get_vertex_count(self):
        """Return amount of vertices in this graph"""
        return len(self.get_vertices())

class SingleVertex(Graph):
    """Special case of a graph. Since we represent a graph as a collection of
    edges, we cannot represent a single-vertex case. This class solves the problem"""
    def __init__(self):
        pass

    def get_vertex_count(self):
        return 1

def count_root_choices(tree, root):
    """Count how many roots (including this one) can be chosen in that
    tree so that rooted tree will be isomorphic or identical to this rooted tree"""

    def get_tuple_representation(tree, root, sans=None):
        p = []
        for v in tree.get_vertices_from(root, sans=sans):
            p.append(get_tuple_representation(tree, v, sans=root))

        return tuple(sorted(p, key=lambda a: len(a)))

    root_tuple = get_tuple_representation(tree, root)

    roots = 1   # because a root can be chosen - we're already picked one (it's root)

    for vert in tree.get_vertices():
        if vert != root:
            if get_tuple_representation(tree, vert) == root_tuple:
                roots += 1

    return roots

def count_labellings_of_rooted_tree(tree, root):
    vertices = tree.get_vertex_count()
    if vertices == 1:  return 1

    tree_kinds = defaultdict(lambda: 0)     # number of vertices in a subtree => number of subtrees

    p = factorial(vertices)

    for subroot, subtree in tree.split_tree(root):
        subvertex_count = subtree.get_vertex_count()
        tree_kinds[subvertex_count] += 1
        p *= count_labellings_of_rooted_tree(subtree, subroot)
        p /= factorial(subvertex_count)

    for dividends in tree_kinds.itervalues():
        p /= factorial(dividends)

    return p


def count_labellings_of_unrooted_tree(tree):
    vertex = list(tree.get_vertices())[0]
    roots = count_root_choices(tree, vertex)
    return count_labellings_of_rooted_tree(tree, vertex) / roots


# Some sample graphs
g1 = Graph({(1,4),(2,4),(3,4),(4,5),(5,6),(5,7),(7,8)})
g2 = Graph({(1, 2), (2, 3)})
g3 = Graph({(1,6),(2,6),(6,5),(5,7),(7,4),(7,3)})
g4 = Graph({(1,2)})
g5 = Graph({(7,4),(7,3),(5,7)})