GraphUtilities.py

from itertools import combinations
from random import sample


def get_combination_two(n):
    return (n * (n - 1)) / 2


def get_combinations(vertex_list):
    temp_combination = combinations(vertex_list, 2)
    combination_list = list()
    for combination in temp_combination:
        combination_list.append(combination)
    return combination_list


def connect_vertices(vertices):
    """Connects vertices with the minimum number of edges... Uses disjoint sets."""
    if len(vertices) < 3:
        return
    parents = [-1 for v in vertices]            # set all parents as -1
    indices = list(range(0, len(vertices)))     # use indices instead of looking for indices every time we sample
    combinations_to_return = list()

    # at this point they are all disjoint
    while parents.count(-1) != 1:               # continue until there is only one -1
        while True:                             # randomly pick two do until we get two disjoint vertices
            p = sample(indices, 2)
            p.sort()                            # sort them to match the combinations generated by python...
            p0 = p[0]
            while parents[p0] != -1:            # find p0's set
                p0 = parents[p0]
            p1 = p[1]
            while parents[p1] != -1:            # find p1's set
                p1 = parents[p1]
            if p0 != p1:                        # if they're not the same then they are disjoint.
                parents[p1] = p[0]              # perform a union operation
                combinations_to_return.append((vertices[p[0]], vertices[p[1]]))
                break
    return combinations_to_return


def get_connected_vertices(vertices, ratio):
    """Connects vertices, and then adds more connections to get the 'ratio'"""
    # the following block connects the clusters
    connected_set = set(connect_vertices(vertices))
    vertex_combinations = set(get_combinations(vertices))
    number_of_samples = (round(len(vertex_combinations) * ratio)) - len(connected_set)
    if number_of_samples < 0:
        number_of_samples = 0
    connected_set = list(set(sample(vertex_combinations.difference(connected_set), number_of_samples)).union(connected_set))
    return connected_set