Random Graph Generator

RandomGraphGenerator.py

from random import *

# Floyd Sampler
def sampler(N, K):
    res = set()
    for i in range(N - K, N):
        j = randint(0, i)
        if j not in res:
            res.add(j)
        else:
            res.add(i)
    return sorted(res)

def edgeToVal(i, j, n):
    return n * (n - 1) // 2 - (n - i - 1) * (n - i) // 2 + j - i - 1

def valToEdge(val, n):
    beg = 0
    end = n - 1
    i = 0
    while beg <= end:
        mid = (beg + end) // 2
        if edgeToVal(mid, mid + 1, n) <= val:
            i = mid
            beg = mid + 1
        else:
            end = mid - 1
    return i, i + 1 + val - edgeToVal(i, i + 1, n)

# N -> Number of vertices
# M -> Number of extra edges to be added to tree
def randGraph(N, M):
    M = min(M, N * (N - 1) // 2 - (N - 1))
    edges = []
    comps = [[i] for i in range(N)]
    for _ in range(N - 1):
        i = randint(0, len(comps) - 1)
        j = randint(0, len(comps) - 2)
        if j == i:
            j += 1
        if len(comps[i]) < len(comps[j]):
            i, j = j, i
        comps[i], comps[-2] = comps[-2], comps[i]
        comps[j], comps[-1] = comps[-1], comps[j]
        u = choice(comps[-2])
        v = choice(comps[-1])
        edges.append((min(u, v), max(u, v)))
        while len(comps[-1]) > 0:
            comps[-2].append(comps[-1].pop())
        comps.pop()
    extraEdges = sampler(N * (N - 1) // 2 - (N - 1), M)
    currentEdgeVals = sorted([edgeToVal(i, j, N) for i, j in edges])
    i = 0
    for k in extraEdges:
        while i < len(currentEdgeVals) and currentEdgeVals[i] - i - 1 < k:
            i += 1
        edges.append(valToEdge(i + k, N))
    return sorted(edges)

for u, v in randGraph(7, 3):
    print(u, v)