Prüfer sequence

prufer_to_tree.py

def prufer_to_tree(A):
    n = len(A)
    G = [[] for i in range(n+2)]
    deg = [1]*(n+2)
    for a in A:
        deg[a] += 1

    E = []
    for a in A:
        for i in range(n+2):
            if deg[i] == 1:
                E.append((a, i))
                deg[a] -= 1
                deg[i] -= 1
                break

    V = [i for i in range(n+2) if deg[i] > 0]
    assert len(V) == 2
    u, v = V
    E.append((u, v))
    return E

def tree_to_prufer(G):
    *D, = map(len, G)
    A = []
    for t in range(N-2):
        for i in range(N):
            if D[i] == 1:
                j, = G[i]
                G[i].remove(j)
                G[j].remove(i)
                D[i] = 0; D[j] -= 1
                A.append(j)
                break
    return A

import random
random.seed()
N = int(input())
seq = [random.randint(0, N-1) for i in range(N-2)]
print(seq)
E = prufer_to_tree(seq)
print(E)

G = [set() for i in range(N)]
for u, v in E:
    G[u].add(v)
    G[v].add(u)

print(tree_to_prufer(G))