Transformation of trees into Prufer sequence and back

prufer_code.py

"""Transformation of trees into Prufer sequence and back.

The Prufer sequence of a labeled tree is a unique seqience associated
with that tree on length n-2 where there are n vertices in the tree.
More information: https://en.wikipedia.org/wiki/Prüfer_sequence

Copyright 2015 Karol M. Langner, Google Inc.
Licensed under the Apache License, Version 2.0
http://www.apache.org/licenses/LICENSE-2.0
"""

def seq2graph(seq, roots, nodes):
  """Turn a sequence into a graph."""

  seq = list(seq)[::-1]
  assert len(seq) == len(nodes) - len(roots)

  assert seq[0] in roots
  dangling = [seq.pop(0)]

  graph = []
  remaining = [r for r in list(nodes) if r not in roots]
  for el in seq:
    if el in remaining:
      graph.append((dangling.pop(0), el))
      remaining.remove(el)
    dangling.append(el)

  assert len(dangling) == len(remaining)
  graph.extend([(d, remaining.pop(0)) for d in dangling])

  return graph

def graph2seq(graph, roots, nodes):
  """Turn a graph into a sequence."""

  roots = list(roots)
  nodes = sorted(list(nodes))
  assert len(graph) == len(nodes) - len(roots)

  parents = [a for a, b in graph]
  leaves = [n for n in nodes if n not in parents]

  seq = []
  while leaves:
    el = leaves.pop(0)

    fathers = [a for a, b in graph if a in parents and b == el]
    assert len(fathers) == 1
    father = fathers[0]

    seq.append(father)
    parents.remove(father)
    if father not in parents and father not in roots:
      leaves.append(father)

  return seq

data = [
  #  seq   roots  nodes    graph
  ( 'AAA', 'A',   'ABCD',  [('A', 'B'), ('A', 'C'), ('A', 'D')] ),
  ( 'CBA', 'A',   'ABCD',  [('A', 'B'), ('B', 'C'), ('C', 'D')] ),
  ( 'BBA', 'A',   'ABCD',  [('A', 'B'), ('B', 'C'), ('B', 'D')] ),
  ( 'CBA', 'AB',  'ABCDE', [('A', 'C'), ('B', 'D'), ('C', 'E')] ),
]
for seq, roots, nodes, graph in data:
  assert seq2graph(seq, roots, nodes) == graph