transform tree

transformedTree.py

# from typing import List, TypeVar

# Tree

#         A
#     /   |   \
#   B    C    D
#  / |    |
# E  F    G

# TNode = TypeVar("TNode", bound="Node")

class Node:

    children = []


# TransformedTree

#      A
#     /
#   B----C----D
#   /     |
#  E--F   G


class Vertex:
    firstChild = None
    nextSibling = None


def generateVertex(node, parentChildren=[]):
    vertex = Vertex()
    if not node.children and not parentChildren:
        return vertex
    if node.children:
        nodeFirstChild = node.children.pop(0)
        vertex.firstChild = generateVertex(nodeFirstChild, node.children)
    if len(parentChildren) >= 1:
        sibling = parentChildren.pop(0)
        vertex.nextSibling = generateVertex(sibling, parentChildren)
    return vertex


def transformTree(root):
    # Node -> Vertex
    v = Vertex()
    children = root.children
    if not children:
        return v
    firstChild = Vertex()
    startChild = children.pop(0)
    v.firstChild = generateVertex(startChild, children)
    if children:
        v.nextSibling = generateVertex(children[0], children)
    return v


if __name__ == '__main__':
    A = Node()
    B = Node()
    C = Node()
    D = Node()
    E = Node()
    F = Node()
    G = Node()

    A.children = [B,C,D]
    B.children = [E,F]
    C.children = [G]

    vertex = Vertex()
    vertex = transformTree(A)
    print vertex.__class__  # A
    print vertex.firstChild.__class__   # B
    print vertex.firstChild.nextSibling.__class__  # C
    print vertex.firstChild.nextSibling.nextSibling.__class__  # D
    print vertex.firstChild.firstChild.nextSibling.__class__  # E
    print vertex.firstChild.nextSibling.nextSibling.__class__  # F
    print vertex.firstChild.nextSibling.firstChild.__class__  # G